import gradio as gr
import random
import os
import torch
import subprocess
import numpy as np
from PIL import Image
from transformers import AutoProcessor, AutoModelForCausalLM
from diffusers import DiffusionPipeline
import cv2
from datetime import datetime
from fastapi import FastAPI

app = FastAPI()


#----------Start of theme----------
theme = gr.themes.Soft(
    primary_hue="zinc",
    secondary_hue="stone",
    font=[gr.themes.GoogleFont('Kavivanar'), gr.themes.GoogleFont('Kavivanar'), 'system-ui', 'sans-serif'],
    font_mono=[gr.themes.GoogleFont('Source Code Pro'), gr.themes.GoogleFont('Inconsolata'), gr.themes.GoogleFont('Inconsolata'), 'monospace'],
).set(
    body_background_fill='*primary_100',
    body_text_color='secondary_600',
    body_text_color_subdued='*primary_500',
    body_text_weight='500',
    background_fill_primary='*primary_100',
    background_fill_secondary='*secondary_200',
    color_accent='*primary_300',
    border_color_accent_subdued='*primary_400',
    border_color_primary='*primary_400',
    block_background_fill='*primary_300',
    block_border_width='*panel_border_width',
    block_info_text_color='*primary_700',
    block_info_text_size='*text_md',
    panel_background_fill='*primary_200',
    accordion_text_color='*primary_600',
    table_text_color='*primary_600',
    input_background_fill='*primary_50',
    input_background_fill_focus='*primary_100',
    button_primary_background_fill='*primary_500',
    button_primary_background_fill_hover='*primary_400',
    button_primary_text_color='*primary_50',
    button_primary_text_color_hover='*primary_100',
    button_cancel_background_fill='*primary_500',
    button_cancel_background_fill_hover='*primary_400'
)
#----------End of theme----------

def flip_image(x):
    return np.fliplr(x)

def basic_filter(image, filter_type):
    """Apply basic image filters"""
    if filter_type == "Gray Toning":
        return cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    elif filter_type == "Sepia":
        sepia_filter = np.array([
            [0.272, 0.534, 0.131],
            [0.349, 0.686, 0.168],
            [0.393, 0.769, 0.189]
        ])
        return cv2.transform(image, sepia_filter)
    elif filter_type == "X-ray":
        # Improved X-ray effect
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        inverted = cv2.bitwise_not(gray)
        # Increase contrast
        clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
        enhanced = clahe.apply(inverted)
        # Sharpen
        kernel = np.array([[-1,-1,-1], [-1,9,-1], [-1,-1,-1]])
        sharpened = cv2.filter2D(enhanced, -1, kernel)
        return cv2.cvtColor(sharpened, cv2.COLOR_GRAY2BGR)
    elif filter_type == "Burn it":
        return cv2.GaussianBlur(image, (15, 15), 0)

def classic_filter(image, filter_type):
    """Classical display filters"""
    if filter_type == "Charcoal Effect":
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        inverted = cv2.bitwise_not(gray)
        blurred = cv2.GaussianBlur(inverted, (21, 21), 0)
        sketch = cv2.divide(gray, cv2.subtract(255, blurred), scale=256)
        return cv2.cvtColor(sketch, cv2.COLOR_GRAY2BGR)
    
    elif filter_type == "Sharpen":
        kernel = np.array([[-1,-1,-1], [-1,9,-1], [-1,-1,-1]])
        return cv2.filter2D(image, -1, kernel)
    
    elif filter_type == "Embossing":
        kernel = np.array([[0,-1,-1], [1,0,-1], [1,1,0]])
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        emboss = cv2.filter2D(gray, -1, kernel) + 128
        return cv2.cvtColor(emboss, cv2.COLOR_GRAY2BGR)
    
    elif filter_type == "Edge Detection":
        edges = cv2.Canny(image, 100, 200)
        return cv2.cvtColor(edges, cv2.COLOR_GRAY2BGR)

def creative_filters(image, filter_type):
    """Creative and unusual image filters"""
    if filter_type == "Pixel Art":
        h, w = image.shape[:2]
        piksel_size = 20
        small = cv2.resize(image, (w//piksel_size, h//piksel_size))
        return cv2.resize(small, (w, h), interpolation=cv2.INTER_NEAREST)
    
    elif filter_type == "Mosaic Effect":
        h, w = image.shape[:2]
        mosaic_size = 30
        for i in range(0, h, mosaic_size):
            for j in range(0, w, mosaic_size):
                roi = image[i:i+mosaic_size, j:j+mosaic_size]
                if roi.size > 0:
                    color = np.mean(roi, axis=(0,1))
                    image[i:i+mosaic_size, j:j+mosaic_size] = color
        return image
    
    elif filter_type == "Rainbow":
        hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        h, w = image.shape[:2]
        for i in range(h):
            hsv[i, :, 0] = (hsv[i, :, 0] + i % 180).astype(np.uint8)
        return cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
    
    elif filter_type == "Night Vision":
        green_image = image.copy()
        green_image[:,:,0] = 0  # Blue channel
        green_image[:,:,2] = 0  # Red channel
        return cv2.addWeighted(green_image, 1.5, np.zeros(image.shape, image.dtype), 0, -50)

def special_effects(image, filter_type):
    """Apply special effects"""
    if filter_type == "Matrix Effect":
        green_matrix = np.zeros_like(image)
        green_matrix[:,:,1] = image[:,:,1]  # Only green channel
        random_brightness = np.random.randint(0, 255, size=image.shape[:2])
        green_matrix[:,:,1] = np.minimum(green_matrix[:,:,1] + random_brightness, 255)
        return green_matrix
    
    elif filter_type == "Wave Effect":
        rows, cols = image.shape[:2]
        img_output = np.zeros(image.shape, dtype=image.dtype)
        
        for i in range(rows):
            for j in range(cols):
                offset_x = int(25.0 * np.sin(2 * 3.14 * i / 180))
                offset_y = int(25.0 * np.cos(2 * 3.14 * j / 180))
                if i+offset_x < rows and j+offset_y < cols:
                    img_output[i,j] = image[(i+offset_x)%rows,(j+offset_y)%cols]
                else:
                    img_output[i,j] = 0
        return img_output
    
    elif filter_type == "Time Stamp":
        output = image.copy()
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        font = cv2.FONT_HERSHEY_SIMPLEX
        cv2.putText(output, timestamp, (10, 30), font, 1, (255, 255, 255), 2)
        return output
    
    elif filter_type == "Glitch Effect":
        glitch = image.copy()
        h, w = image.shape[:2]
        for _ in range(10):
            x1 = random.randint(0, w-50)
            y1 = random.randint(0, h-50)
            x2 = random.randint(x1, min(x1+50, w))
            y2 = random.randint(y1, min(y1+50, h))
            glitch[y1:y2, x1:x2] = np.roll(glitch[y1:y2, x1:x2], 
                                          random.randint(-20, 20), 
                                          axis=random.randint(0, 1))
        return glitch

def artistic_filters(image, filter_type):
    """Applies artistic image filters"""
    if filter_type == "Pop Art":
        img_small = cv2.resize(image, None, fx=0.5, fy=0.5)
        img_color = cv2.resize(img_small, (image.shape[1], image.shape[0]))
        for _ in range(2):
            img_color = cv2.bilateralFilter(img_color, 9, 300, 300)
        hsv = cv2.cvtColor(img_color, cv2.COLOR_BGR2HSV)
        hsv[:,:,1] = hsv[:,:,1]*1.5
        return cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
    
    elif filter_type == "Oil Paint":
        ret = np.float32(image.copy())
        ret = cv2.bilateralFilter(ret, 9, 75, 75)
        ret = cv2.detailEnhance(ret, sigma_s=15, sigma_r=0.15)
        ret = cv2.edgePreservingFilter(ret, flags=1, sigma_s=60, sigma_r=0.4)
        return np.uint8(ret)
    
    elif filter_type == "Cartoon":
        # Improved cartoon effect
        color = image.copy()
        gray = cv2.cvtColor(color, cv2.COLOR_BGR2GRAY)
        gray = cv2.medianBlur(gray, 5)
        edges = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 9, 9)
        color = cv2.bilateralFilter(color, 9, 300, 300)
        cartoon = cv2.bitwise_and(color, color, mask=edges)
        # Increase color saturation
        hsv = cv2.cvtColor(cartoon, cv2.COLOR_BGR2HSV)
        hsv[:,:,1] = hsv[:,:,1]*1.4  # saturation increase
        return cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)

def atmospheric_filters(image, filter_type):
    """atmospheric filters"""
    if filter_type == "Autumn":
        # Genhanced autumn effect
        autumn_filter = np.array([
            [0.393, 0.769, 0.189],
            [0.349, 0.686, 0.168],
            [0.272, 0.534, 0.131]
        ])
        autumn = cv2.transform(image, autumn_filter)
        # Increase color temperature
        hsv = cv2.cvtColor(autumn, cv2.COLOR_BGR2HSV)
        hsv[:,:,0] = hsv[:,:,0]*0.8  # Shift to orange/yellow tones
        hsv[:,:,1] = hsv[:,:,1]*1.2  # Increase saturation
        return cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
    
    elif filter_type == "Nostalgia":
        # Improved nostalgia effect
        # Reduce contrast and add yellowish tone
        image = cv2.convertScaleAbs(image, alpha=0.9, beta=10)
        sepia = cv2.transform(image, np.array([
            [0.393, 0.769, 0.189],
            [0.349, 0.686, 0.168],
            [0.272, 0.534, 0.131]
        ]))
        # Darkening effect in corners
        h, w = image.shape[:2]
        kernel = np.zeros((h, w))
        center = (h//2, w//2)
        for i in range(h):
            for j in range(w):
                dist = np.sqrt((i-center[0])**2 + (j-center[1])**2)
                kernel[i,j] = 1 - min(1, dist/(np.sqrt(h**2 + w**2)/2))
        kernel = np.dstack([kernel]*3)
        return cv2.multiply(sepia, kernel).astype(np.uint8)
    
    elif filter_type == "Increase Brightness":
        # Improved brightness boost
        hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        # Increase brightness
        hsv[:,:,2] = cv2.convertScaleAbs(hsv[:,:,2], alpha=1.2, beta=30)
        # Also increase the contrast slightly
        return cv2.convertScaleAbs(cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR), alpha=1.1, beta=0)

def image_processing(image, filter_type):
    """Main image processing function"""
    if image is None:
        return None
    
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    
    # Process by filter categories
    basic_filter_list = ["Gray Toning", "Sepia", "X-ray", "Burn it"]
    classic_filter_list = ["Charcoal Effect", "Sharpen", "Embossing", "Edge Detection"]
    creative_filters_list = ["Rainbow", "Night Vision"]
    special_effects_list = ["Matrix Effect", "Wave Effect", "Time Stamp", "Glitch Effect"]
    artistic_filters_list = ["Pop Art", "Oil Paint", "Cartoon"]
    atmospheric_filters_list = ["Autumn", "Increase Brightness"]
    
    if filter_type in basic_filter_list:
        output = basic_filter(image, filter_type)
    elif filter_type in classic_filter_list:
        output = classic_filter(image, filter_type)
    elif filter_type in creative_filters_list:
        output = creative_filters(image, filter_type)
    elif filter_type in special_effects_list:
        output = special_effects(image, filter_type)
    elif filter_type in artistic_filters_list:
        output = artistic_filters(image, filter_type)
    elif filter_type in atmospheric_filters_list:
        output = atmospheric_filters(image, filter_type)
    else:
        output = image
        
    return cv2.cvtColor(output, cv2.COLOR_BGR2RGB) if len(output.shape) == 3 else output

    # Get absolute path of image file
image_path = 'https://huggingface.co/spaces/DigiP-AI/Image_Studio/blob/main/abstract.jpg' # Replace with your image file path

absolute_path = os.path.abspath(image_path)
    
css = """
.gradio-container {
    background: url(https://huggingface.co/spaces/DigiP-AI/Image_Studio/blob/main/abstract.jpg)
}
"""

# Gradio interface
with gr.Blocks(theme=theme, css=css) as app:
    gr.HTML("<center><h6>🎨 Image Studio</h6></center>")

    with gr.Tab("Image to Prompt"): 
        subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)

        # Initialize Florence model
        device = "cuda" if torch.cuda.is_available() else "cpu"
        florence_model = AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True).to(device).eval()
        florence_processor = AutoProcessor.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True)

        # api_key = os.getenv("HF_READ_TOKEN")
        
        def generate_caption(image):
            if not isinstance(image, Image.Image):
                image = Image.fromarray(image)
            
            inputs = florence_processor(text="<MORE_DETAILED_CAPTION>", images=image, return_tensors="pt").to(device)
            generated_ids = florence_model.generate(
                input_ids=inputs["input_ids"],
                pixel_values=inputs["pixel_values"],
                max_new_tokens=1024,
                early_stopping=False,
                do_sample=False,
                num_beams=3,
            )
            generated_text = florence_processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
            parsed_answer = florence_processor.post_process_generation(
                generated_text,
                task="<MORE_DETAILED_CAPTION>",
                image_size=(image.width, image.height)
            )
            prompt =  parsed_answer["<MORE_DETAILED_CAPTION>"]
            print("\n\nGeneration completed!:"+ prompt)
            return prompt

        io = gr.Interface(generate_caption,
            inputs=[gr.Image(label="Input Image")],
            outputs = [gr.Textbox(label="Output Prompt", lines=2, show_copy_button = True),
            # gr.Image(label="Output Image")
                       ]
                    )
    
    with gr.Tab("Text to Image"):
        gr.HTML("<center><h6>ℹ️ Please do not run the models at the same time, the models are currently running on the CPU, which might affect performance.</h6></center>")        
        with gr.Accordion("Flux-RealismLora", open=False):
            model1 = gr.load("models/XLabs-AI/flux-RealismLora")
        with gr.Accordion("Flux--schnell-realism", open=False):
            model2 = gr.load("models/hugovntr/flux-schnell-realism")
        with gr.Accordion("Flux--schnell-LoRA", open=False):
            model3 = gr.load("models/Octree/flux-schnell-lora")
        
    with gr.Tab("Flip Image"):
                with gr.Row():
                    image_input = gr.Image(type="numpy", label="Upload Image")
                    image_output = gr.Image(format="png")
                with gr.Row():    
                    image_button = gr.Button("Run", variant='primary')
                    image_button.click(flip_image, inputs=image_input, outputs=image_output)
    with gr.Tab("Image Filters"):    
        with gr.Row():
            with gr.Column():
                image_input = gr.Image(type="numpy", label="Upload Image")
                with gr.Accordion("ℹ️ Filter Categories", open=True):
                    filter_type = gr.Dropdown(
                        [
                            # Basic Filters
                            "Gray Toning", "Sepia", "X-ray", "Burn it",
                            # Classic Filter 
                            "Charcoal Effect", "Sharpen", "Embossing", "Edge Detection",                       
                            # Creative Filters
                            "Rainbow", "Night Vision",
                            # Special Effects
                            "Matrix Effect", "Wave Effect", "Time Stamp", "Glitch Effect",
                            # Artistic Filters
                            "Pop Art", "Oil Paint", "Cartoon",
                            # Atmospheric Filters
                            "Autumn", "Increase Brightness"
                        ],
                        label="🎭 Select Filter",
                        info="Choose the effect you want"
                    )
                submit_button = gr.Button("✨ Apply Filter", variant="primary")
    
            with gr.Column():
                image_output = gr.Image(label="🖼️ Filtered Image")
                
            submit_button.click(
                image_processing,
                inputs=[image_input, filter_type],
                outputs=image_output
            )

    
            
    with gr.Tab("Image Upscaler"):    
        with gr.Row():
            with gr.Column():
                def upscale_image(input_image, radio_input):
                    upscale_factor = radio_input
                    output_image = cv2.resize(input_image, None, fx = upscale_factor, fy = upscale_factor, interpolation = cv2.INTER_CUBIC)
                    return output_image
                
                radio_input = gr.Radio(label="Upscale Levels", choices=[2, 4, 6, 8, 10], value=2)
                
                iface = gr.Interface(fn=upscale_image, inputs = [gr.Image(label="Input Image", interactive=True), radio_input], outputs = gr.Image(label="Upscaled Image", format="png"), title="Image Upscaler")
    
app.launch(share=True)