from __future__ import annotations

import numpy as np

from .chroma import BackgroundModel
from .settings import DespillSettings


def despill_green(
    rgb: np.ndarray,
    alpha: np.ndarray,
    model: BackgroundModel,
    settings: DespillSettings,
    screen_color: tuple[float, float, float] | None = None,
) -> tuple[np.ndarray, np.ndarray]:
    rgb_f = rgb.astype(np.float32) / 255.0
    if not settings.enabled:
        return rgb.copy(), np.zeros(alpha.shape, dtype=np.float32)
    if screen_color is not None:
        # Arbitrary background colour: despill along its Lab chroma direction.
        return _despill_color(rgb, alpha, model, settings)

    # (1) Despill across the whole foreground, not only the soft edge band.
    # Green frequently bleeds into pixels ViTMatte marks as solid foreground
    # (alpha ~ 1); the old upper bound left those uncorrected. Neutral pixels
    # stay untouched because the green-excess term below is zero for them.
    foreground = alpha > settings.edge_low
    if settings.edge_expand_radius > 0:
        from .deps import require_cv2

        cv2 = require_cv2()
        radius = settings.edge_expand_radius
        kernel = cv2.getStructuringElement(
            cv2.MORPH_ELLIPSE, (radius * 2 + 1, radius * 2 + 1)
        )
        foreground = cv2.dilate(foreground.astype(np.uint8), kernel).astype(bool)

    r = rgb_f[..., 0]
    g = rgb_f[..., 1]
    b = rgb_f[..., 2]
    neutral_green = np.maximum(r, b) + settings.green_excess_margin
    excess = np.maximum(g - neutral_green, 0.0)

    # (2) Keep full strength across the whole opaque range (alpha >= 0.5) and
    # ease off only where the pixel is mostly transparent. The previous tent
    # peaked at alpha 0.5 and starved exactly the solid hair that needs despill.
    alpha_weight = np.maximum(
        np.clip(alpha * 2.0, 0.0, 1.0), settings.alpha_weight_floor
    )

    # (3) Localise purely by green excess. bg_confidence is intentionally not a
    # gate: it is low on the character edge where the spill lives, so using it
    # suppressed the very correction we want.
    green_weight = np.clip(excess / max(settings.green_excess_margin * 8.0, 1e-6), 0.0, 1.0)
    mask = foreground.astype(np.float32) * alpha_weight * green_weight

    out = rgb_f.copy()
    target_green = np.minimum(g, neutral_green)
    out[..., 1] = g * (1.0 - mask * settings.strength) + target_green * (
        mask * settings.strength
    )

    # Light foreground color approximation: green spill usually removed too much
    # luminance from thin hair edges, so redistribute part of the removed green
    # into red/blue without changing fully opaque pixels.
    bg_green = float(model.rgb_center[1])
    compensation = excess * mask * settings.strength * min(0.25, bg_green * 0.15)
    out[..., 0] = np.clip(out[..., 0] + compensation * 0.5, 0.0, 1.0)
    out[..., 2] = np.clip(out[..., 2] + compensation * 0.5, 0.0, 1.0)

    return np.clip(out * 255.0, 0, 255).astype(np.uint8), mask.astype(np.float32)


def _despill_color(
    rgb: np.ndarray,
    alpha: np.ndarray,
    model: BackgroundModel,
    settings: DespillSettings,
) -> tuple[np.ndarray, np.ndarray]:
    """Despill toward an arbitrary background colour.

    Generalises the green channel algebra: the background's Lab a*/b* vector gives
    a chroma direction; foreground chroma projecting along it beyond a margin is
    treated as spill and pulled back, keeping luminance so texture survives. Reduces
    to a green-style suppression when the background is a saturated primary, but also
    handles low-chroma pastels where no single RGB channel dominates.
    """
    from .deps import require_cv2

    cv2 = require_cv2()
    rgb_f = rgb.astype(np.float32) / 255.0

    bg_lab = np.asarray(model.lab_center, dtype=np.float32)
    va, vb = float(bg_lab[1]), float(bg_lab[2])
    mag = float(np.hypot(va, vb))
    if mag < 1e-3:
        return rgb.copy(), np.zeros(alpha.shape, dtype=np.float32)
    da, db = va / mag, vb / mag

    foreground = alpha > settings.edge_low
    if settings.edge_expand_radius > 0:
        radius = settings.edge_expand_radius
        kernel = cv2.getStructuringElement(
            cv2.MORPH_ELLIPSE, (radius * 2 + 1, radius * 2 + 1)
        )
        foreground = cv2.dilate(foreground.astype(np.uint8), kernel).astype(bool)

    lab = cv2.cvtColor(rgb_f, cv2.COLOR_RGB2LAB)
    A, B = lab[..., 1], lab[..., 2]
    proj = A * da + B * db                       # chroma toward the background hue
    excess = np.maximum(proj - settings.color_margin_lab, 0.0)

    alpha_weight = np.maximum(
        np.clip(alpha * 2.0, 0.0, 1.0), settings.alpha_weight_floor
    )
    color_weight = np.clip(
        excess / max(settings.color_margin_lab * 8.0, 1e-6), 0.0, 1.0
    )
    mask = foreground.astype(np.float32) * alpha_weight * color_weight
    reduce = excess * mask * settings.strength   # pull chroma back along the bg axis

    lab_out = lab.copy()
    lab_out[..., 1] = A - reduce * da
    lab_out[..., 2] = B - reduce * db
    out = cv2.cvtColor(lab_out, cv2.COLOR_LAB2RGB)
    return np.clip(out * 255.0, 0, 255).astype(np.uint8), mask.astype(np.float32)