# app.py

import gradio as gr
import numpy as np
import torch
from PIL import Image, ImageFilter
from torchvision import transforms
from transformers import AutoModelForImageSegmentation


DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
MODEL_ID = "ZhengPeng7/BiRefNet"

model = AutoModelForImageSegmentation.from_pretrained(
    MODEL_ID,
    trust_remote_code=True
).to(DEVICE).float()

model.eval()

transform = transforms.Compose([
    transforms.Resize((1024, 1024)),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    ),
])


def predict_subject_mask(image: Image.Image) -> Image.Image:
    """Returns white subject / black background mask."""
    image_rgb = image.convert("RGB")
    original_size = image_rgb.size

    x = transform(image_rgb).unsqueeze(0).to(DEVICE).float()

    with torch.no_grad():
        output = model(x)

        # BiRefNet often returns: (scaled_preds, class_preds)
        if isinstance(output, (list, tuple)):
            if isinstance(output[0], (list, tuple)):
                pred = output[0][-1]
            else:
                pred = output[0]
        elif hasattr(output, "logits"):
            pred = output.logits
        else:
            pred = output

        pred = pred.sigmoid().cpu()[0].squeeze()

    mask = transforms.ToPILImage()(pred)
    mask = mask.resize(original_size, Image.BILINEAR)
    return mask


def lock_background(
    reference,
    target,
    threshold=128,
    feather=8,
    protect_both_frames=True,
):
    ref = Image.fromarray(reference).convert("RGB")
    tgt = Image.fromarray(target).convert("RGB")

    if ref.size != tgt.size:
        tgt = tgt.resize(ref.size, Image.BICUBIC)

    ref_np = np.array(ref).astype(np.float32)
    tgt_np = np.array(tgt).astype(np.float32)

    ref_mask = predict_subject_mask(ref)
    tgt_mask = predict_subject_mask(tgt)

    ref_m = np.array(ref_mask)
    tgt_m = np.array(tgt_mask)

    # foreground = white, background = black
    ref_bg = ref_m < threshold
    tgt_bg = tgt_m < threshold

    if protect_both_frames:
        shared_bg = ref_bg & tgt_bg
    else:
        shared_bg = tgt_bg

    alpha = Image.fromarray((shared_bg * 255).astype(np.uint8))
    alpha = alpha.filter(ImageFilter.GaussianBlur(radius=feather))
    alpha_np = np.array(alpha).astype(np.float32) / 255.0
    alpha_np = alpha_np[..., None]

    # copy canonical background from reference into target
    fixed_np = tgt_np * (1 - alpha_np) + ref_np * alpha_np
    fixed_np = np.clip(fixed_np, 0, 255).astype(np.uint8)

    # diagnostic difference before/after
    before_diff = np.abs(ref_np - tgt_np).mean(axis=2)
    after_diff = np.abs(ref_np - fixed_np.astype(np.float32)).mean(axis=2)

    heatmap = np.clip((before_diff - after_diff) * 4, 0, 255).astype(np.uint8)
    heatmap_rgb = np.stack([heatmap, np.zeros_like(heatmap), 255 - heatmap], axis=2)

    shared_bg_vis = np.stack([
        shared_bg.astype(np.uint8) * 255,
        shared_bg.astype(np.uint8) * 255,
        shared_bg.astype(np.uint8) * 255,
    ], axis=2)

    before_score = float(before_diff[shared_bg].mean()) if shared_bg.any() else 0
    after_score = float(after_diff[shared_bg].mean()) if shared_bg.any() else 0

    report = f"""
### Background Lock Report

- Shared background pixels: **{shared_bg.mean() * 100:.2f}%**
- Mean background difference before: **{before_score:.2f}**
- Mean background difference after: **{after_score:.2f}**
- Improvement: **{before_score - after_score:.2f}**
"""

    return fixed_np, shared_bg_vis, heatmap_rgb, report


demo = gr.Interface(
    fn=lock_background,
    inputs=[
        gr.Image(label="Reference frame", type="numpy"),
        gr.Image(label="End frame", type="numpy"),
        gr.Slider(0, 255, value=128, step=1, label="Foreground protection threshold"),
        gr.Slider(0, 40, value=0, step=1, label="Feather edge"),
        gr.Checkbox(value=True, label="Protect foreground in both frames"),
    ],
    outputs=[
        gr.Image(label="Fixed end frame", format="png"),
        gr.Image(label="Shared background mask", format="png"),
        gr.Image(label="Improvement heatmap", format="png"),
        gr.Markdown(label="Report"),
    ],
    examples=[
        ["Frame1.jpg", "Frame2.jpg", 32, 0, True],
        ["Frame2.jpg", "Frame1.jpg", 32, 0, True],
    ],
    title="End Frame Background Lock",
    description=(
        "Locks the background of an end frame to a reference frame for aligned, "
        "locked-off AI-generated shots while preserving the foreground subject."
    ),
)

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