"""Crop-around-person and composite-back helpers for person replacement.

Uses the SAM3 pose mask to find the person bbox, crops the driving video
around them (so the generation model's centered output aligns), and pastes
the generated person back into the original frames.
"""

import cv2
import numpy as np
import torch


def person_bbox(mask, thresh=0.1):
    """Union bbox of the person across all mask frames.

    mask: [B, H, W, 3] colored mask (person = any bright channel).
    Returns (x0, y0, x1, y1) or None if nothing found.
    """
    m = (mask[..., :3].max(dim=-1).values > thresh)  # B,H,W bool
    ys, xs = torch.nonzero(m.any(dim=0), as_tuple=True)
    if len(xs) == 0:
        return None
    return int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1


def fit_crop_box(bbox, frame_w, frame_h, target_ar, margin=0.15):
    """Expand bbox by margin, adjust to target aspect ratio, clamp to frame."""
    x0, y0, x1, y1 = bbox
    bw, bh = x1 - x0, y1 - y0
    cx, cy = (x0 + x1) / 2, (y0 + y1) / 2
    bw *= 1 + margin
    bh *= 1 + margin
    # fit aspect
    if bw / bh > target_ar:
        bh = bw / target_ar
    else:
        bw = bh * target_ar
    # clamp size to frame
    bw = min(bw, frame_w)
    bh = min(bh, frame_h)
    x0 = int(round(cx - bw / 2))
    y0 = int(round(cy - bh / 2))
    x0 = max(0, min(x0, frame_w - int(bw)))
    y0 = max(0, min(y0, frame_h - int(bh)))
    return x0, y0, x0 + int(bw), y0 + int(bh)


def crop_video(video, box):
    """Crop ComfyUI IMAGE tensor [B,H,W,C] to (x0,y0,x1,y1)."""
    x0, y0, x1, y1 = box
    return video[:, y0:y1, x0:x1, :]


def resize_video(video, target_w, target_h):
    """Resize ComfyUI IMAGE tensor [B,H,W,C] with lanczos via cv2."""
    arr = (video.cpu().numpy() * 255).astype(np.uint8)
    out = np.stack([
        cv2.resize(f, (target_w, target_h), interpolation=cv2.INTER_LANCZOS4)
        for f in arr
    ])
    return torch.from_numpy(out).float() / 255.0


def composite_back(original, generated, box, mask=None, feather=8):
    """Paste generated person back into original frames.

    original:  [B, H, W, 3] float 0-1 (source video)
    generated: [B, h, w, 3] float 0-1 (generated, same AR as box)
    box:       (x0, y0, x1, y1) paste location
    mask:      optional [B, h, w] float 0-1 person alpha (None = paste all)
    """
    x0, y0, x1, y1 = box
    bw, bh = x1 - x0, y1 - y0
    gen = resize_video(generated, bw, bh)
    out = original.clone()
    if mask is None:
        out[:, y0:y1, x0:x1, :] = gen
    else:
        alpha = mask.cpu().numpy()
        if feather > 0:
            alpha = np.stack([
                cv2.GaussianBlur(a, (feather * 2 + 1,) * 2, 0) for a in alpha
            ])
        alpha_t = torch.from_numpy(alpha).float().unsqueeze(-1)  # B,h,w,1
        region = out[:, y0:y1, x0:x1, :]
        out[:, y0:y1, x0:x1, :] = gen * alpha_t + region * (1 - alpha_t)
    return out