import torch

# --------------------------------------------------------------------
#  Blend mode helpers – kept as module-level functions for clean reuse
# --------------------------------------------------------------------
def blend_normal(bg_rgb, fg_rgb):
    return fg_rgb

def blend_screen(bg_rgb, fg_rgb):
    return 1.0 - (1.0 - bg_rgb) * (1.0 - fg_rgb)

def blend_multiply(bg_rgb, fg_rgb):
    return bg_rgb * fg_rgb

def blend_overlay(bg_rgb, fg_rgb):
    result = torch.where(bg_rgb <= 0.5, 2.0 * bg_rgb * fg_rgb,
                         1.0 - 2.0 * (1.0 - bg_rgb) * (1.0 - fg_rgb))
    return result.clamp(0.0, 1.0)

def blend_difference(bg_rgb, fg_rgb):
    return torch.abs(bg_rgb - fg_rgb)

def blend_dissolve(bg_rgb, fg_rgb):
    noise = torch.rand_like(bg_rgb)
    return torch.where(noise < 0.5, fg_rgb, bg_rgb)

def blend_clear(bg_rgb, fg_rgb):
    return bg_rgb

def blend_add(bg_rgb, fg_rgb):
    return torch.clamp(bg_rgb + fg_rgb, 0.0, 1.0)

def blend_subtract(bg_rgb, fg_rgb):
    return torch.clamp(bg_rgb - fg_rgb, 0.0, 1.0)

def blend_replace(bg_rgb, fg_rgb):
    return fg_rgb

def blend_darken_only(bg_rgb, fg_rgb):
    return torch.min(bg_rgb, fg_rgb)

def blend_lighten_only(bg_rgb, fg_rgb):
    return torch.max(bg_rgb, fg_rgb)

def blend_color_dodge(bg_rgb, fg_rgb):
    result = bg_rgb / (1.0 - fg_rgb + 1e-8)
    return result.clamp(0.0, 1.0)

def blend_color_burn(bg_rgb, fg_rgb):
    result = 1.0 - (1.0 - bg_rgb) / (fg_rgb + 1e-8)
    return result.clamp(0.0, 1.0)

def blend_hard_light(bg_rgb, fg_rgb):
    result = torch.where(fg_rgb <= 0.5, 2.0 * bg_rgb * fg_rgb,
                         1.0 - 2.0 * (1.0 - bg_rgb) * (1.0 - fg_rgb))
    return result.clamp(0.0, 1.0)

def blend_soft_light(bg_rgb, fg_rgb):
    result = torch.where(fg_rgb <= 0.5,
                         (2.0 * bg_rgb * fg_rgb) + (bg_rgb * bg_rgb * (1.0 - 2.0 * fg_rgb)),
                         2.0 * bg_rgb * (1.0 - fg_rgb) + torch.sqrt(bg_rgb) * (2.0 * fg_rgb - 1.0))
    return result.clamp(0.0, 1.0)

def blend_pin_light(bg_rgb, fg_rgb):
    result = torch.where(fg_rgb <= 0.5,
                         torch.min(bg_rgb, 2.0 * fg_rgb),
                         torch.max(bg_rgb, 2.0 * fg_rgb - 1.0))
    return result.clamp(0.0, 1.0)

def blend_linear_dodge_add(bg_rgb, fg_rgb):
    return torch.clamp(bg_rgb + fg_rgb, 0.0, 1.0)

def blend_linear_burn_multiply(bg_rgb, fg_rgb):
    return torch.clamp(bg_rgb + fg_rgb - 1.0, 0.0, 1.0)

def blend_vivid_light(bg_rgb, fg_rgb):
    result = torch.where(fg_rgb <= 0.5,
                         1.0 - (1.0 - bg_rgb) / (2.0 * fg_rgb + 1e-8),
                         bg_rgb / (2.0 * (1.0 - fg_rgb) + 1e-8))
    return result.clamp(0.0, 1.0)

def blend_halo(bg_rgb, fg_rgb):
    edge = torch.abs(bg_rgb - fg_rgb)
    glow = (bg_rgb + fg_rgb) * edge
    return glow.clamp(0.0, 1.0)

def blend_negate(bg_rgb, fg_rgb):
    return 1.0 - torch.abs(bg_rgb - fg_rgb)

def blend_merge_rgb(bg_rgb, fg_rgb):
    return fg_rgb

def blend_merge_cyan(bg_rgb, fg_rgb):
    cyan = 1.0 - fg_rgb[..., 0:1]
    return cyan.repeat(1, 1, 1, 3)

def blend_merge_magenta(bg_rgb, fg_rgb):
    magenta = 1.0 - fg_rgb[..., 1:2]
    return magenta.repeat(1, 1, 1, 3)

def blend_merge_yellow(bg_rgb, fg_rgb):
    yellow = 1.0 - fg_rgb[..., 2:3]
    return yellow.repeat(1, 1, 1, 3)

def blend_merge_black(bg_rgb, fg_rgb):
    intensity = fg_rgb.mean(dim=-1, keepdim=True)
    return intensity.repeat(1, 1, 1, 3)

def blend_fill_fg_color(bg_rgb, fg_rgb):
    return fg_rgb

def blend_fill_bg_color(bg_rgb, fg_rgb):
    return bg_rgb

def blend_average(bg_rgb, fg_rgb):
    return (bg_rgb + fg_rgb) * 0.5

def blend_composite(bg_rgb, fg_rgb):
    return fg_rgb

def blend_behind(bg_rgb, fg_rgb):
    return fg_rgb

def blend_template(bg_rgb, fg_rgb):
    return fg_rgb

def composite_layer_stack(background, layer, blending_mode, transparency, mask=None):
    """
    background, layer: tensors of shape [B, H, W, C] where C=3 or 4
    mask: optional tensor [B, H, W] or [H, W]; values in [0,1]
    transparency: global opacity multiplier
    """
    device = background.device
    layer = layer.to(device)
    bg_h, bg_w = background.shape[1:3]

    # ---------- 1. Ensure background and layer have 4 channels ----------
    bg = background.clone()
    if bg.shape[-1] == 3:
        alpha_bg = torch.ones((bg.shape[0], bg_h, bg_w, 1), device=device)
        bg = torch.cat([bg, alpha_bg], dim=-1)
    else:
        alpha_bg = bg[..., 3:4]

    fg = layer.clone()
    if fg.shape[-1] == 3:
        alpha_fg = torch.ones((fg.shape[0], fg.shape[1], fg.shape[2], 1), device=device)
        fg = torch.cat([fg, alpha_fg], dim=-1)
    else:
        alpha_fg = fg[..., 3:4]

    bg_rgb = bg[..., :3]
    fg_rgb = fg[..., :3]

    # ---------- 2. Resize layer to match background dimensions ----------
    if fg.shape[1:3] != (bg_h, bg_w):
        fg_rgb = torch.nn.functional.interpolate(
            fg_rgb.permute(0, 3, 1, 2), size=(bg_h, bg_w), mode='bilinear', align_corners=False
        ).permute(0, 2, 3, 1)
        alpha_fg = torch.nn.functional.interpolate(
            alpha_fg.permute(0, 3, 1, 2), size=(bg_h, bg_w), mode='bilinear', align_corners=False
        ).permute(0, 2, 3, 1)

    # ---------- 3. Prepare mask ----------
    if mask is None:
        mask = torch.ones((bg.shape[0], bg_h, bg_w, 1), device=device)
    else:
        mask = mask.to(device)
        if mask.dim() == 2:          # (H, W)
            mask = mask.unsqueeze(0).unsqueeze(-1)
        elif mask.dim() == 3:        # (B, H, W) or (B, H, W, 1)
            if mask.shape[-1] != 1:
                mask = mask.unsqueeze(-1)
        if mask.shape[1:3] != (bg_h, bg_w):
            mask = torch.nn.functional.interpolate(
                mask.permute(0, 3, 1, 2), size=(bg_h, bg_w), mode='bilinear', align_corners=False
            ).permute(0, 2, 3, 1)
        if mask.shape[0] < bg.shape[0]:
            mask = mask.expand(bg.shape[0], -1, -1, -1)

    # ---------- 4. Compute final opacity factor ----------
    opacity = alpha_fg * mask * transparency   # [B, H, W, 1]

    # Special handling for "Clear" and "Behind"
    if blending_mode.lower() == "clear":
        opacity = torch.zeros_like(opacity)
    elif blending_mode.lower() == "behind":
        opacity = opacity * (1.0 - alpha_bg)

    # ---------- 5. Apply RGB blend mode ----------
    # Build a mapping from the existing module-level blend functions
    blend_func_name = blending_mode.lower().replace(' ', '_').replace('(', '').replace(')', '').replace('-', '_')
    # Use the functions defined at module level (they are in the same scope)
    blend_func = globals().get(f"blend_{blend_func_name}", blend_normal)

    blended_rgb = blend_func(bg_rgb, fg_rgb)

    # ---------- 6. Composite RGB using opacity ----------
    result_rgb = bg_rgb * (1.0 - opacity) + blended_rgb * opacity

    # ---------- 7. Compute result alpha (standard over compositing) ----------
    result_alpha = alpha_bg + (1.0 - alpha_bg) * opacity

    # ---------- 8. Concatenate and clamp ----------
    result = torch.cat([result_rgb, result_alpha], dim=-1).clamp(0.0, 1.0)

    return (result,)