import os
import cv2
import numpy as np
import torch
import gradio as gr
import spaces
from gradio.themes.base import Base
from gradio.themes.utils import colors, fonts, sizes
from PIL import Image, ImageOps
from transformers import AutoModelForImageSegmentation
from torchvision import transforms

class WhiteTheme(Base):
    def __init__(
        self,
        *,
        primary_hue: colors.Color | str = colors.orange,
        font: fonts.Font | str | tuple[fonts.Font | str, ...] = (
            fonts.GoogleFont("Inter"),
            "ui-sans-serif",
            "system-ui",
            "sans-serif",
        ),
        font_mono: fonts.Font | str | tuple[fonts.Font | str, ...] = (
            fonts.GoogleFont("Inter"),
            "ui-monospace",
            "system-ui",
            "monospace",
        )
    ):
        super().__init__(
            primary_hue=primary_hue,
            font=font,
            font_mono=font_mono,
        )
        
        self.set(
            # Light mode specific colors
            background_fill_primary="*primary_50",
            background_fill_secondary="white",
            border_color_primary="*primary_300",
            
            # General colors that should stay constant
            body_background_fill="white",
            body_background_fill_dark="white",
            block_background_fill="white",
            block_background_fill_dark="white",
            panel_background_fill="white",
            panel_background_fill_dark="white",
            body_text_color="black",
            body_text_color_dark="black",
            block_label_text_color="black",
            block_label_text_color_dark="black",
            block_border_color="white",
            panel_border_color="white",
            input_border_color="lightgray",
            input_background_fill="white",
            input_background_fill_dark="white",
            shadow_drop="none"
        )

torch.set_float32_matmul_precision('high')
torch.jit.script = lambda f: f

device = "cuda" if torch.cuda.is_available() else "cpu"

def refine_foreground(image, mask, r=90):
    if mask.size != image.size:
        mask = mask.resize(image.size)
    image = np.array(image) / 255.0
    mask = np.array(mask) / 255.0
    estimated_foreground = FB_blur_fusion_foreground_estimator_2(image, mask, r=r)
    image_masked = Image.fromarray((estimated_foreground * 255.0).astype(np.uint8))
    return image_masked

def FB_blur_fusion_foreground_estimator_2(image, alpha, r=90):
    alpha = alpha[:, :, None]
    F, blur_B = FB_blur_fusion_foreground_estimator(
        image, image, image, alpha, r)
    return FB_blur_fusion_foreground_estimator(image, F, blur_B, alpha, r=6)[0]

def FB_blur_fusion_foreground_estimator(image, F, B, alpha, r=90):
    if isinstance(image, Image.Image):
        image = np.array(image) / 255.0
    blurred_alpha = cv2.blur(alpha, (r, r))[:, :, None]
    blurred_FA = cv2.blur(F * alpha, (r, r))
    blurred_F = blurred_FA / (blurred_alpha + 1e-5)
    blurred_B1A = cv2.blur(B * (1 - alpha), (r, r))
    blurred_B = blurred_B1A / ((1 - blurred_alpha) + 1e-5)
    F = blurred_F + alpha * (image - alpha * blurred_F - (1 - alpha) * blurred_B)
    F = np.clip(F, 0, 1)
    return F, blurred_B

class ImagePreprocessor():
    def __init__(self, resolution=(1024, 1024)) -> None:
        self.transform_image = transforms.Compose([
            transforms.Resize(resolution),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406],
                                 [0.229, 0.224, 0.225]),
        ])

    def proc(self, image: Image.Image) -> torch.Tensor:
        image = self.transform_image(image)
        return image

# Load the model
birefnet = AutoModelForImageSegmentation.from_pretrained(
    'zhengpeng7/BiRefNet-matting', trust_remote_code=True)
birefnet.to(device)
birefnet.eval()

def remove_background_wrapper(image):
    if image is None:
        raise gr.Error("Please upload an image.")
    image_ori = Image.fromarray(image).convert('RGB')
    foreground, background, pred_pil, reverse_mask = remove_background(image_ori)
    return foreground, background, pred_pil, reverse_mask

@spaces.GPU
def remove_background(image_ori):
    original_size = image_ori.size
    image_preprocessor = ImagePreprocessor(resolution=(1024, 1024))
    image_proc = image_preprocessor.proc(image_ori)
    image_proc = image_proc.unsqueeze(0)
    
    with torch.no_grad():
        preds = birefnet(image_proc.to(device))[-1].sigmoid().cpu()
    pred = preds[0].squeeze()
    
    pred_pil = transforms.ToPILImage()(pred)
    pred_pil = pred_pil.resize(original_size, Image.BICUBIC)
    
    reverse_mask = ImageOps.invert(pred_pil)
    
    foreground = image_ori.copy()
    foreground.putalpha(pred_pil)
    
    background = image_ori.copy()
    background.putalpha(reverse_mask)
    
    torch.cuda.empty_cache()
    
    return foreground, background, pred_pil, reverse_mask

# Custom CSS for styling
custom_css = """
.title-container {
    text-align: center;
    padding: 10px 0;
}

#title {
    font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Helvetica, Arial, sans-serif;
    font-size: 36px;
    font-weight: bold;
    color: #000000;
    padding: 10px;
    border-radius: 10px;
    display: inline-block;
    background: linear-gradient(
        135deg,
        #e0f7fa, #e8f5e9, #fff9c4, #ffebee,
        #f3e5f5, #e1f5fe, #fff3e0, #e8eaf6
    );
    background-size: 400% 400%;
    animation: gradient-animation 15s ease infinite;
}

@keyframes gradient-animation {
    0% { background-position: 0% 50%; }
    50% { background-position: 100% 50%; }
    100% { background-position: 0% 50%; }
}

#submit-button {
    background: linear-gradient(
        135deg,
        #e0f7fa, #e8f5e9, #fff9c4, #ffebee,
        #f3e5f5, #e1f5fe, #fff3e0, #e8eaf6
    );
    background-size: 400% 400%;
    animation: gradient-animation 15s ease infinite;
    border-radius: 12px;
    color: black;
}

/* Force light mode styles */
:root, :root[data-theme='light'], :root[data-theme='dark'] {
    --body-background-fill: white !important;
    --background-fill-primary: white !important;
    --background-fill-secondary: white !important;
    --block-background-fill: white !important;
    --panel-background-fill: white !important;
    --body-text-color: black !important;
    --block-label-text-color: black !important;
}

/* Additional overrides for dark mode */
@media (prefers-color-scheme: dark) {
    :root {
        color-scheme: light;
    }
}
"""

with gr.Blocks(css=custom_css, theme=WhiteTheme()) as demo:
    gr.HTML('''
        <div class="title-container">
            <div id="title">
                <span>{.</span><span id="typed-text"></span><span>}</span>
            </div>
        </div>
        <script>
            (function() {
                const text = "image";
                const typedTextSpan = document.getElementById("typed-text");
                let charIndex = 0;
                
                function type() {
                    if (charIndex < text.length) {
                        typedTextSpan.textContent += text[charIndex];
                        charIndex++;
                        setTimeout(type, 150);
                    }
                }
                
                setTimeout(type, 150);
            })();
        </script>
    ''')
    
    # Interface setup with input and output
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="numpy", sources=['upload'], label="Upload Image")
            btn = gr.Button("Process Image", elem_id="submit-button")
        
        with gr.Column():
            output_foreground = gr.Image(type="pil", label="Foreground")
            output_background = gr.Image(type="pil", label="Background")
            output_foreground_mask = gr.Image(type="pil", label="Foreground Mask")
            output_background_mask = gr.Image(type="pil", label="Background Mask")

    # Link the button to the processing function
    btn.click(fn=remove_background_wrapper, inputs=image_input, outputs=[
        output_foreground, output_background, output_foreground_mask, output_background_mask])

    demo.launch(debug=True)