import numpy as np
import torch
import torch.nn.functional as F
from torchvision.transforms.functional import normalize
import gradio as gr
from briarmbg import BriaRMBG
import PIL
from PIL import Image
import tempfile
import os
import time
import uuid
import shutil

# Load the pre-trained model
print("Loading model...")
net = BriaRMBG.from_pretrained("briaai/RMBG-1.4")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net.to(device)
net.eval()
print(f"Model loaded on {device}")

# Create output directory if it doesn't exist
OUTPUT_DIR = "output_images"
os.makedirs(OUTPUT_DIR, exist_ok=True)

def process(image, progress=gr.Progress()):
    if image is None:
        return None, None, None
    try:
        progress(0, desc="Starting processing...")
        orig_image = Image.fromarray(image)
        original_size = orig_image.size
        
        progress(0.2, desc="Preparing image...")
        process_image = orig_image.resize(original_size, Image.LANCZOS)
        w, h = process_image.size
        
        im_np = np.array(process_image)
        im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2, 0, 1)
        im_tensor = torch.unsqueeze(im_tensor, 0)
        im_tensor = torch.divide(im_tensor, 255.0)
        im_tensor = normalize(im_tensor, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
        
        progress(0.4, desc="Processing with AI model...")
        if torch.cuda.is_available():
            im_tensor = im_tensor.cuda()
        
        with torch.no_grad():
            result = net(im_tensor)
        
        progress(0.6, desc="Post-processing...")
        result = torch.squeeze(F.interpolate(result[0][0], size=(h, w), mode='bilinear'), 0)
        ma = torch.max(result)
        mi = torch.min(result)
        result = (result - mi) / (ma - mi)
        
        result_array = (result * 255).cpu().data.numpy().astype(np.uint8)
        pil_mask = Image.fromarray(np.squeeze(result_array))
        
        if pil_mask.size != original_size:
            pil_mask = pil_mask.resize(original_size, Image.LANCZOS)
        
        new_im = orig_image.copy()
        new_im.putalpha(pil_mask)
        
        progress(0.8, desc="Saving result...")
        unique_id = str(uuid.uuid4())[:8]
        filename = f"background_removed_{unique_id}.png"
        filepath = os.path.join(OUTPUT_DIR, filename)
        new_im.save(filepath, format='PNG', quality=100)
        
        # Convert to numpy array for display
        output_array = np.array(new_im.convert('RGBA'))
        
        progress(1.0, desc="Done!")
        return (
            output_array, 
            gr.update(value=filepath, visible=True),
            gr.update(value=f"""
                <script>
                    setTimeout(function() {{
                        window.location.href = '/file={filepath}';
                    }}, 1000);
                </script>
            """)
        )
    
    except Exception as e:
        print(f"Error processing image: {str(e)}")
        return None, None, None

css = """
@import url('https://fonts.googleapis.com/css2?family=Orbitron:wght@400;500;700&display=swap');

.title-text {
    color: #ff00de;
    font-family: 'Orbitron', sans-serif;
    font-size: 2.5em;
    text-align: center;
    margin: 20px 0;
    text-shadow: 0 0 10px rgba(255, 0, 222, 0.7);
    animation: glow 2s ease-in-out infinite alternate;
}

.subtitle-text {
    color: #00ffff;
    text-align: center;
    margin-bottom: 30px;
    font-size: 1.2em;
    text-shadow: 0 0 8px rgba(0, 255, 255, 0.7);
}

.image-container {
    background: rgba(10, 10, 30, 0.3);
    border-radius: 15px;
    padding: 20px;
    margin: 10px 0;
    border: 2px solid #00ffff;
    box-shadow: 0 0 15px rgba(0, 255, 255, 0.2);
    transition: all 0.3s ease;
}

.image-container img {
    max-width: 100%;
    height: auto;
    display: block;
    margin: 0 auto;
}

.image-container:hover {
    box-shadow: 0 0 20px rgba(0, 255, 255, 0.4);
    transform: translateY(-2px);
}

.download-btn {
    background: linear-gradient(45deg, #00ffff, #ff00de);
    border: none;
    padding: 12px 25px;
    border-radius: 8px;
    color: white;
    font-family: 'Orbitron', sans-serif;
    cursor: pointer;
    transition: all 0.3s ease;
    margin-top: 10px;
    text-align: center;
    text-transform: uppercase;
    letter-spacing: 1px;
    width: 100%;
    display: block;
}

.download-btn:hover {
    transform: translateY(-2px);
    box-shadow: 0 5px 15px rgba(0, 255, 255, 0.4);
}

@keyframes glow {
    from { text-shadow: 0 0 5px #ff00de, 0 0 10px #ff00de; }
    to { text-shadow: 0 0 10px #ff00de, 0 0 20px #ff00de; }
}

@media (max-width: 768px) {
    .title-text { font-size: 1.8em; }
    .subtitle-text { font-size: 1em; }
    .image-container { padding: 10px; }
    .download-btn { padding: 10px 20px; }
}
"""

with gr.Blocks(css=css) as demo:
    gr.Markdown("""
        <h1 class="title-text">AI Background Removal</h1>
        <p class="subtitle-text">Remove backgrounds instantly using advanced AI technology</p>
    """)
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(
                label="Upload Image",
                type="numpy",
                elem_classes="image-container"
            )
            
            output_image = gr.Image(
                label="Result",
                type="numpy",
                show_label=True,
                elem_classes="image-container"
            )
            
            download_button = gr.File(
                label="Download Processed Image",
                visible=True,
                elem_classes="download-btn"
            )
            
            auto_download = gr.HTML(visible=False)
    
    input_image.change(
        fn=process,
        inputs=input_image,
        outputs=[output_image, download_button, auto_download]
    )

if __name__ == "__main__":
    demo.launch()