import os
import cv2
import numpy as np
import torch
import gradio as gr
import spaces  # Required for @spaces.GPU

from PIL import Image, ImageOps
from transformers import AutoModelForImageSegmentation
from torchvision import transforms

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')
    # Call the processing function
    foreground, background, pred_pil, reverse_mask = remove_background(image_ori)
    return foreground, background, pred_pil, reverse_mask

@spaces.GPU  # Decorate the processing function
def remove_background(image_ori):
    original_size = image_ori.size

    # Preprocess the image
    image_preprocessor = ImagePreprocessor(resolution=(1024, 1024))
    image_proc = image_preprocessor.proc(image_ori)
    image_proc = image_proc.unsqueeze(0)

    # Prediction
    with torch.no_grad():
        preds = birefnet(image_proc.to(device))[-1].sigmoid().cpu()
    pred = preds[0].squeeze()

    # Process Results
    pred_pil = transforms.ToPILImage()(pred)
    pred_pil = pred_pil.resize(original_size, Image.BICUBIC)  # Resize mask to original size

    # Create reverse mask (background mask)
    reverse_mask = ImageOps.invert(pred_pil)

    # Create foreground image (object with transparent background)
    foreground = image_ori.copy()
    foreground.putalpha(pred_pil)

    # Create background image
    background = image_ori.copy()
    background.putalpha(reverse_mask)

    torch.cuda.empty_cache()

    # Return images in the specified order
    return foreground, background, pred_pil, reverse_mask

# Custom CSS for button styling
custom_css = """
@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;
}
"""

with gr.Blocks(css=custom_css) as demo:
    # Interface setup with input and output
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="numpy", 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)