"""Person-aware background blur: composite, hold, I420, encode path."""
import numpy as np

from portrait import (
    MaskHold,
    apply_i420,
    apply_portrait,
    bgr_to_i420,
    composite_bgr,
    i420_nbytes,
    i420_source_path,
    i420_to_bgr,
    person_present,
    soft_blur_bgr,
)


def _checker(h: int, w: int) -> np.ndarray:
    img = np.zeros((h, w, 3), dtype=np.uint8)
    img[::4, :] = 255
    img[:, ::4] = 255
    return img


def test_composite_keeps_person_pixels():
    h, w = 80, 128
    bg = _checker(h, w)
    person = bg.copy()
    person[16:64, 32:96] = (40, 180, 40)
    mask = np.zeros((h, w), dtype=np.float32)
    mask[16:64, 32:96] = 1.0
    out = composite_bgr(person, soft_blur_bgr(person, 9), mask)
    roi = out[24:56, 40:88]
    assert abs(int(roi[:, :, 1].mean()) - 180) < 8
    assert abs(int(roi[:, :, 0].mean()) - 40) < 8


def test_composite_blurs_background_checker():
    h, w = 80, 128
    bg = _checker(h, w)
    person = bg.copy()
    person[16:64, 32:96] = (40, 180, 40)
    mask = np.zeros((h, w), dtype=np.float32)
    mask[16:64, 32:96] = 1.0
    sharp_var = float(bg[0:12, :].var())
    out = composite_bgr(person, soft_blur_bgr(person, 9), mask)
    blur_var = float(out[0:12, :].var())
    assert sharp_var > 1000
    assert blur_var < sharp_var * 0.5


def test_empty_mask_is_fully_blurred():
    img = _checker(40, 64)
    mask = np.zeros((40, 64), dtype=np.float32)
    out = apply_portrait(img, mask)
    assert person_present(mask) is False
    assert float(out.var()) < float(img.var()) * 0.5
    assert not np.array_equal(out, img)


def test_tiny_blob_is_not_a_person():
    mask = np.zeros((40, 64), dtype=np.float32)
    mask[10:12, 10:12] = 0.9
    assert person_present(mask) is False


def test_mask_hold_keeps_last_on_miss():
    hold = MaskHold(hold_s=0.5, ema=1.0)
    mask = np.ones((8, 8), dtype=np.float32)
    got = hold.update(mask, True, 1.0)
    assert got is not None
    missed = hold.update(None, False, 1.2)
    assert missed is not None
    assert float(missed.mean()) == 1.0
    assert hold.update(None, False, 1.6) is None


def test_ema_smooths_mask_flicker():
    hold = MaskHold(hold_s=1.0, ema=0.5)
    a = np.zeros((4, 4), dtype=np.float32)
    b = np.ones((4, 4), dtype=np.float32)
    hold.update(a, True, 0.0)
    out = hold.update(b, True, 0.1)
    assert 0.4 < float(out.mean()) < 0.6


def test_i420_roundtrip_keeps_640x360():
    w, h = 640, 360
    bgr = np.zeros((h, w, 3), dtype=np.uint8)
    bgr[20:200, 80:400] = (30, 90, 200)
    payload = bgr_to_i420(bgr)
    assert len(payload) == i420_nbytes(w, h)
    mask = np.zeros((h, w), dtype=np.float32)
    mask[20:200, 80:400] = 1.0
    out = apply_i420(payload, w, h, mask, radius=9)
    assert len(out) == i420_nbytes(w, h)
    back = i420_to_bgr(out, w, h)
    assert back.shape == (h, w, 3)


def test_missing_mask_blurs_whole_frame():
    w, h = 48, 32
    bgr = _checker(h, w)
    payload = bgr_to_i420(bgr)
    out = apply_i420(payload, w, h, None, radius=7)
    y0 = np.frombuffer(payload, dtype=np.uint8, count=w * h).reshape(h, w)
    y1 = np.frombuffer(out, dtype=np.uint8, count=w * h).reshape(h, w)
    assert float(y1.var()) < float(y0.var()) * 0.6


def test_apply_i420_blurs_y_background():
    w, h = 64, 48
    bgr = np.zeros((h, w, 3), dtype=np.uint8)
    bgr[::3, :] = 255
    bgr[:, ::3] = 255
    bgr[8:40, 16:48] = (20, 200, 20)
    payload = bgr_to_i420(bgr)
    mask = np.zeros((h, w), dtype=np.float32)
    mask[8:40, 16:48] = 1.0
    out = apply_i420(payload, w, h, mask, radius=9)
    y0 = np.frombuffer(payload, dtype=np.uint8, count=w * h).reshape(h, w)
    y1 = np.frombuffer(out, dtype=np.uint8, count=w * h).reshape(h, w)
    assert float(y1[0:6, :].var()) < float(y0[0:6, :].var()) * 0.6
    assert abs(int(y1[12:36, 20:44].mean()) - int(y0[12:36, 20:44].mean())) < 12


def test_c920_uses_portrait_dir_rally_stays_raw():
    assert i420_source_path("cam-01", "/run/raw", "/run/por") == "/run/por/cam-01.i420"
    assert i420_source_path("cam-15", "/run/raw", "/run/por") == "/run/por/cam-15.i420"
    assert i420_source_path("rally", "/run/raw", "/run/por") == "/run/raw/rally.i420"
    assert i420_source_path("cam-01", "/run/raw", "") == "/run/raw/cam-01.i420"


def test_letterbox_mask_maps_back_to_16x9():
    from portrait import letterbox_bgr, unletterbox_mask

    bgr = np.zeros((360, 640, 3), dtype=np.uint8)
    canvas, box = letterbox_bgr(bgr, 256)
    assert canvas.shape == (256, 256, 3)
    x0, y0, nw, nh = box
    assert nw > nh
    m = np.zeros((256, 256), dtype=np.float32)
    cy, cx = y0 + nh // 2, x0 + nw // 2
    m[cy - 2:cy + 3, cx - 2:cx + 3] = 1.0
    full = unletterbox_mask(m, box, 360, 640)
    assert full.shape == (360, 640)
    assert float(full[180, 320]) > 0.5
    assert float(full[2, 2]) < 0.1


def test_selfie_cpu_sees_warrior_not_empty():
    import cv2
    from pathlib import Path

    from portrait import PersonSeg

    rgb = np.load(Path(__file__).parent / "fixtures" / "movenet_warrior_192.npy")
    bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
    seg = PersonSeg()
    assert seg.device == "CPU"
    mask = seg.infer_bgr(bgr)
    assert person_present(mask)
    empty = np.full((360, 640, 3), 40, dtype=np.uint8)
    assert person_present(seg.infer_bgr(empty)) is False


def test_pose_protect_keeps_wrists_and_face():
    from portrait import pose_protect_mask

    k = np.zeros((17, 3), dtype=np.float32)
    k[0] = [0.30, 0.50, 0.9]   # nose
    k[9] = [0.12, 0.18, 0.9]   # L wrist
    k[10] = [0.12, 0.82, 0.9]  # R wrist
    m = pose_protect_mask(100, 200, k)
    assert float(m[12, 36]) > 0.7
    assert float(m[12, 164]) > 0.7
    assert float(m[30, 100]) > 0.7
    assert float(m[90, 10]) < 0.2


def test_union_masks_keeps_face_and_body():
    from portrait import union_masks

    body = np.zeros((40, 40), dtype=np.float32)
    body[20:36, 10:30] = 1.0
    face = np.zeros((40, 40), dtype=np.float32)
    face[4:14, 14:26] = 1.0
    out = union_masks(body, face)
    assert float(out[8, 20]) == 1.0
    assert float(out[28, 20]) == 1.0
    assert float(out[2, 2]) == 0.0


def test_lighter_blur_keeps_more_background_detail():
    w, h = 64, 48
    bgr = np.zeros((h, w, 3), dtype=np.uint8)
    bgr[::3, :] = 255
    bgr[:, ::3] = 255
    payload = bgr_to_i420(bgr)
    mask = np.zeros((h, w), dtype=np.float32)
    mask[8:40, 16:48] = 1.0
    heavy = apply_i420(payload, w, h, mask, radius=15)
    light = apply_i420(payload, w, h, mask, radius=7)
    y0 = np.frombuffer(payload, dtype=np.uint8, count=w * h).reshape(h, w)
    yh = np.frombuffer(heavy, dtype=np.uint8, count=w * h).reshape(h, w)
    yl = np.frombuffer(light, dtype=np.uint8, count=w * h).reshape(h, w)
    bg0 = float(y0[0:6, :].var())
    assert float(yl[0:6, :].var()) > float(yh[0:6, :].var())
    assert float(yl[0:6, :].var()) < bg0