#!/usr/bin/env python3
"""One-off black-matte recovery for the erase-v5 blush comparison."""

from pathlib import Path

import numpy as np
from PIL import Image


HERE = Path(__file__).resolve().parent
SOURCE = HERE.parent / "split-v2" / "raw" / "blush-v2.png"
OUTPUT = HERE / "parts" / "blush-unmatted-v1.png"
NOISE_FLOOR = 6


raw = np.asarray(Image.open(SOURCE).convert("RGB"), dtype=np.float32)
coverage = raw.max(axis=2)
keep = coverage >= NOISE_FLOOR

alpha = np.where(keep, coverage, 0.0)
straight = np.zeros_like(raw)
straight[keep] = np.clip(raw[keep] * (255.0 / coverage[keep, None]), 0.0, 255.0)

rgba = np.dstack((np.rint(straight), np.rint(alpha))).astype(np.uint8)
Image.fromarray(rgba, "RGBA").save(OUTPUT)

strong = coverage >= 64
ratio_before = raw[strong] / np.maximum(raw[strong].sum(axis=1, keepdims=True), 1.0)
ratio_after = straight[strong] / np.maximum(straight[strong].sum(axis=1, keepdims=True), 1.0)
print({
    "output": str(OUTPUT),
    "noiseFloor": NOISE_FLOOR,
    "nonzeroAlpha": int(np.count_nonzero(alpha)),
    "alphaAbove64": int(np.count_nonzero(alpha > 64)),
    "darkRgbWithAlphaAbove64": int(np.count_nonzero((straight.max(axis=2) < 80) & (alpha > 64))),
    "strongPinkMaxChromaticityDelta": float(np.max(np.abs(ratio_before - ratio_after))),
})