"""
subtitle_mask.py
===================
VapourSynth subtitle mask creation, based on InpaintDelogo.avsi (v3.7) by VoodooFX.
Original: https://github.com/Purfview/InpaintDelogo

Implements the dynamic mask generation paths used for subtitle detection:
  - DynMask=2 : Combo mask  — DynTune (luma binarise) + optional DynColor + DynTEdge
  - DynMask=4 : Subtitle-with-halo mask — DynTune + DynMask4H hysteresis
  - DynMask=6 : Luma-only binarise mask (DynTune)

The returned clip is a single-plane (GRAY8/GRAY16) binary mask in the same
spatial region as the source (full frame or a cropped 'loc' area).

Required plugins
----------------
  - vapoursynth-descale / resize  (built-in resizers)
  - vs-placebo or similar for colour-key masking   [DynColor path, optional]
  - akarin / std.Expr             (built-in Expr) – akarin is preferred if available
  - rgvs (RemoveGrain)            [KillNoise paths]
  - tedgemask                     [DynTEdge path, optional]
"""

from __future__ import annotations
import vapoursynth as vs
from vapoursynth import core
from typing import Union, Optional

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _to_gray(clip: vs.VideoNode) -> vs.VideoNode:
    """
    Extract luma as GRAY8, preserving the original range.

    We do NOT force range expansion here — luma values stay as-is so that
    thresholds in _binarise/_dyntune_mask mean the same thing as what you
    see in a frame inspector (e.g. 16-235 for limited, 0-255 for full).
    """
    if clip.format.color_family == vs.GRAY:
        if clip.format.id != vs.GRAY8:
            clip = core.resize.Bilinear(clip, format=vs.GRAY8)
        return clip
    # ShufflePlanes pulls the Y plane without any range/matrix conversion,
    # so the luma values are exactly what they are in the source.
    return core.std.ShufflePlanes(clip, planes=0, colorfamily=vs.GRAY)


def _binarise(clip: vs.VideoNode, threshold: int) -> vs.VideoNode:
    """Full-range GRAY8 binarise: pixels >= threshold → 255, else 0."""
    assert clip.format.id == vs.GRAY8
    return core.std.Binarize(clip, threshold=threshold)


def _inflate(clip: vs.VideoNode, radius: int = 1) -> vs.VideoNode:
    """Morphological dilation (maximum filter) by *radius* iterations."""
    for _ in range(radius):
        clip = core.std.Maximum(clip)
    return clip


def _inpand(clip: vs.VideoNode, radius: int = 1) -> vs.VideoNode:
    """Morphological erosion (minimum filter) by *radius* iterations."""
    for _ in range(radius):
        clip = core.std.Minimum(clip)
    return clip


# ---------------------------------------------------------------------------
# Plugin capability detection — resolved once at import time
# ---------------------------------------------------------------------------

def _has_plugin(namespace: str, func: str) -> bool:
    return hasattr(core, namespace) and hasattr(getattr(core, namespace), func)

_USE_ZSMOOTH    = _has_plugin("zsmooth",    "RemoveGrain")
_USE_MISC_HYS   = _has_plugin("misc",        "Hysteresis")   # vs-miscfilters-obsolete
_USE_HYS_HYS    = _has_plugin("hysteresis",  "Hysteresis")   # sgt0/vapoursynth-hysteresis
_USE_AKARIN_EXPR = _has_plugin("akarin",     "Expr")

print(f"[subtitle_mask] hysteresis.Hysteresis={_USE_HYS_HYS}  "
      f"misc.Hysteresis={_USE_MISC_HYS}  "
      f"zsmooth.RemoveGrain={_USE_ZSMOOTH}  "
      f"akarin.Expr={_USE_AKARIN_EXPR}")

# ---------------------------------------------------------------------------
# Expression wrapper – uses akarin.Expr if available, otherwise std.Expr
# ---------------------------------------------------------------------------

def _expr(clips, expr):
    """
    Wrapper for core.akarin.Expr (preferred) or core.std.Expr.

    Automatically converts a single clip into a list for consistency.
    """
    if not isinstance(clips, list):
        clips = [clips]
    if _USE_AKARIN_EXPR:
        return core.akarin.Expr(clips, expr)
    else:
        return core.std.Expr(clips, expr)


def _remove_grain(clip: vs.VideoNode, mode: int = 2) -> vs.VideoNode:
    """
    RemoveGrain mode wrapper.

    Prefers zsmooth.RemoveGrain when available — it is a modern,
    maintained reimplementation with identical mode semantics and is
    generally faster.  Falls back to rgvs.RemoveGrain otherwise.

    https://github.com/adworacz/zsmooth
    """
    if _USE_ZSMOOTH:
        return core.zsmooth.RemoveGrain(clip, mode=mode)
    return core.rgvs.RemoveGrain(clip, mode=mode)


def _hysteresis(clip_weak: vs.VideoNode,
                clip_strong: vs.VideoNode) -> vs.VideoNode:
    """
    Hysteresis: grow clip_strong into clip_weak.

    Both plugins share the same argument order:
      clipa = strong seed  (the smaller, reliable mask)
      clipb = weak candidate  (the larger region to grow into)

    Priority: hysteresis.Hysteresis (sgt0) → misc.Hysteresis → AND fallback.
    The AND fallback is a last resort — it only keeps pixels where both
    masks directly overlap, with no flood-fill, so mode 4 will be poor.
    """
    if _USE_HYS_HYS:
        return core.hysteresis.Hysteresis(clip_strong, clip_weak)
    if _USE_MISC_HYS:
        return core.misc.Hysteresis(clip_strong, clip_weak)
    # Fallback: simple bitwise AND (no flood-fill, much less accurate)
    return _expr([clip_weak, clip_strong], "x y and")


# ---------------------------------------------------------------------------
# Diagnostics
# ---------------------------------------------------------------------------

def debug_levels(clip: vs.VideoNode,
                 loc: Optional[tuple[int, int, int, int]] = None,
                 frame: int = 0) -> dict:
    """
    Print the min/max/average luma in the subtitle region for a given frame.
    Use this to find the right dyn_tune and dynmask4h thresholds.

    Example
    -------
    import subtitle_mask
    info = subtitle_mask.debug_levels(src, loc=(100,1380,60,440), frame=500)
    # Then set dyn_tune just below the max, dynmask4h just above the min.
    """
    if loc is not None:
        left, top, right, bottom = loc
        clip = core.std.Crop(clip, left=left, top=top,
                             right=right, bottom=bottom)

    gray = core.std.ShufflePlanes(clip, planes=0, colorfamily=vs.GRAY) \
           if clip.format.color_family != vs.GRAY else clip

    f = gray.get_frame(frame)
    import array
    # Read raw plane data - memoryview is the clean VS way
    plane = f[0]
    data = array.array('B', bytes(plane))

    mn  = min(data)
    mx  = max(data)
    avg = sum(data) / len(data)

    # Bucket histogram: 16 buckets of 16 values each (0-15, 16-31, ... 240-255)
    buckets = [0] * 16
    for v in data:
        buckets[v >> 4] += 1
    total = len(data)

    print(f"[debug_levels] frame={frame}  min={mn}  max={mx}  avg={avg:.1f}  total_px={total}")
    print(f"  Luma histogram (each bar = % of pixels in that 16-value bucket):")
    for i, count in enumerate(buckets):
        lo = i * 16
        hi = lo + 15
        pct = 100 * count / total
        bar = "█" * int(pct / 2)
        print(f"    {lo:3d}-{hi:3d}  {bar:25s} {pct:5.1f}%  ({count} px)")

    # Find the top-2 peaks to guess text and outline luma ranges
    sorted_b = sorted(enumerate(buckets), key=lambda x: -x[1])
    peak1_idx = sorted_b[0][0]
    peak2_idx = sorted_b[1][0]
    peaks = sorted([peak1_idx, peak2_idx])
    text_lo  = peaks[1] * 16        # higher-luma peak = text
    text_hi  = text_lo + 15
    dark_lo  = peaks[0] * 16        # lower-luma peak = outline/background
    dark_hi  = dark_lo + 15

    print(f"  Two dominant luma bands: {dark_lo}-{dark_hi} and {text_lo}-{text_hi}")
    print(f"  → Suggested dyn_tune  (text lower bound):   {text_lo}")
    print(f"  → Suggested dynmask4h (outline upper bound): {dark_hi}")
    print(f"  Note: for coloured subs (yellow etc.) luma is lower than white.")
    print(f"        Use dyn_color for reliable colour-based detection.")

    result = {"min": mn, "max": mx, "avg": round(avg, 1), "frame": frame,
              "histogram": buckets,
              "suggested_dyn_tune": text_lo,
              "suggested_dynmask4h": dark_hi}
    return result


def debug_color(clip: vs.VideoNode,
                loc: Optional[tuple[int, int, int, int]] = None,
                frame: int = 0,
                dyn_tune: int = 140) -> None:
    """
    Sample the average RGB colour of pixels above dyn_tune threshold.
    Use this to find the dyn_color hex value for coloured subtitles.

    Example
    -------
    subtitle_mask.debug_color(src, loc=(100,1380,60,440), frame=150, dyn_tune=140)
    # Output: Avg colour of bright pixels: R=255 G=255 B=0  →  dyn_color=0xFFFF00
    """
    if loc is not None:
        left, top, right, bottom = loc
        clip = core.std.Crop(clip, left=left, top=top,
                             right=right, bottom=bottom)

    # Convert to RGB24 for easy R/G/B sampling
    rgb = core.resize.Bilinear(clip, format=vs.RGB24, matrix_in_s="709")
    gray = core.std.ShufflePlanes(clip, planes=0, colorfamily=vs.GRAY)            if clip.format.color_family != vs.GRAY else clip

    import array
    fg = gray.get_frame(frame)
    fr = rgb.get_frame(frame)
    y_data = array.array('B', bytes(fg[0]))
    # VS RGB24 plane order: 0=G, 1=B, 2=R
    g_data = array.array('B', bytes(fr[0]))
    b_data = array.array('B', bytes(fr[1]))
    r_data = array.array('B', bytes(fr[2]))

    # Sample only pixels above the luma threshold (likely text pixels)
    r_vals, g_vals, b_vals = [], [], []
    for i, y in enumerate(y_data):
        if y >= dyn_tune:
            r_vals.append(r_data[i])
            g_vals.append(g_data[i])
            b_vals.append(b_data[i])

    if not r_vals:
        print(f"[debug_color] No pixels above dyn_tune={dyn_tune}. Try lowering it.")
        return

    avg_r = int(sum(r_vals) / len(r_vals))
    avg_g = int(sum(g_vals) / len(g_vals))
    avg_b = int(sum(b_vals) / len(b_vals))
    hex_color = (avg_r << 16) | (avg_g << 8) | avg_b
    print(f"[debug_color] frame={frame}  sampled {len(r_vals)} pixels above Y>={dyn_tune}")
    print(f"  Avg colour: R={avg_r}  G={avg_g}  B={avg_b}  →  dyn_color=0x{hex_color:06X}")
    print(f"  Use: subtitle_mask(clip, dyn_mask_mode=2, dyn_color=0x{hex_color:06X}, "
          f"dyn_color_tol=20, dyn_tune={dyn_tune})")


def debug_mask_stages(clip: vs.VideoNode,
                      loc: Optional[tuple[int, int, int, int]] = None,
                      dyn_tune: int = 200,
                      dynmask4h: int = 60,
                      frame: int = 0) -> None:
    """
    Print pixel statistics for each mask stage to diagnose an all-black result.

    Example
    -------
    subtitle_mask.debug_mask_stages(src, loc=(100,1380,60,440),
                                    dyn_tune=200, dynmask4h=60, frame=500)
    """
    if loc is not None:
        left, top, right, bottom = loc
        clip = core.std.Crop(clip, left=left, top=top,
                             right=right, bottom=bottom)

    gray = core.std.ShufflePlanes(clip, planes=0, colorfamily=vs.GRAY)            if clip.format.color_family != vs.GRAY else clip

    def _stats(name: str, node: vs.VideoNode) -> None:
        import array
        f = node.get_frame(frame)
        data = array.array('B', bytes(f[0]))
        white = sum(1 for v in data if v > 0)
        total = len(data)
        print(f"  {name:30s}  min={min(data):3d}  max={max(data):3d}"
              f"  white_px={white}/{total} ({100*white//total}%)")

    print(f"\n[debug_mask_stages] frame={frame}  "
          f"clip={clip.width}x{clip.height}  "
          f"format={clip.format.name}")
    print(f"  hysteresis backend: "
          f"{'hysteresis' if _USE_HYS_HYS else 'misc' if _USE_MISC_HYS else 'AND-fallback (inaccurate — install hysteresis or misc plugin!)'}")
    _stats("luma (raw Y plane)",       gray)
    _stats(f"dyn_tune>={dyn_tune}",    core.std.Binarize(gray, threshold=dyn_tune))
    _stats(f"dynmask4h<={dynmask4h}",  _expr(gray, f"x {dynmask4h} <= 255 0 ?"))

    text_mask = core.std.Binarize(gray, threshold=dyn_tune)
    halo_mask = _expr(gray, f"x {dynmask4h} <= 255 0 ?")
    _stats("halo_mask (dark pixels)",  halo_mask)

    # Find minimum inflate needed for text seed to touch the dark outline
    print("  Inflate sweep (text seed vs halo overlap):")
    for inf in range(0, 6):
        seed = text_mask
        for _ in range(inf):
            seed = core.std.Maximum(seed)
        overlap = _expr([seed, halo_mask], "x y and")
        import array as _arr
        f2 = overlap.get_frame(frame)
        d2 = _arr.array('B', bytes(f2[0]))
        white = sum(1 for v in d2 if v > 0)
        print(f"    inflate={inf}  overlap white_px={white}"
              + ("  ← first hit" if inf > 0 and white > 0 else "")
              + ("  ← use this or higher" if white > 0 else "  (no overlap yet)"))
        if white > 0 and inf >= 2:
            break
    print("  → If all inflates show 0 overlap, dyn_tune is too high or dynmask4h too low.")
    print("    Try lowering dyn_tune by 20 or raising dynmask4h by 20 and re-run.")


# ---------------------------------------------------------------------------
# KillNoise — mirrors InpaintDelogo's KillNoise parameter
# ---------------------------------------------------------------------------

def _kill_noise(clip: vs.VideoNode, mode: int) -> vs.VideoNode:
    """
    Reduce small mask artifacts.

    mode 0 : disabled
    mode 1 : single Undot  (RemoveGrain mode 2)
    mode 2 : two-pass Undot
    mode 3 : Hysteresis cleanup  (inpand → hysteresis)
    mode 4 : Undot + Hysteresis
    mode 5 : two-pass Undot + Hysteresis

    Negative modes pre-apply expand→inpand to reduce sponginess.
    """
    if mode == 0:
        return clip

    if mode < 0:
        # Pre-sponge reduction
        clip = _inflate(clip)
        clip = _inpand(clip)
        mode = -mode

    if mode in (1, 4, 5):
        clip = _remove_grain(clip, mode=2)
    if mode in (2, 5):
        clip = _remove_grain(clip, mode=2)
    if mode in (3, 4, 5):
        eroded = _inpand(clip)
        clip = _hysteresis(clip, eroded)

    return clip


# ---------------------------------------------------------------------------
# KillBlobs — mirrors InpaintDelogo's KillBlobs parameter
# ---------------------------------------------------------------------------

def _kill_blobs(clip: vs.VideoNode,
                strength: int,
                pre_blobs: int = 0) -> vs.VideoNode:
    """
    Remove blobs thicker than logo/subtitle text.

    strength > 0 : inpand→expand method  (safer, leaves blob edges)
    strength < 0 : inpand→hysteresis method  (removes blobs fully)
    pre_blobs    : pre-sponginess reduction (±1 = 'both' mode, ±2 = 'square')

    Mirrors AviSynth's positive/negative KillBlobs logic.
    """
    if strength == 0:
        return clip

    n = abs(strength)

    # Pre-blob sponginess reduction
    if pre_blobs != 0:
        if pre_blobs > 0:
            expanded = _inflate(clip)
            clip = _kill_blobs_inner(clip, n, strength > 0)
            clip = _inpand(expanded)  # not quite right but close
        else:
            expanded = _inflate(clip)
            eroded   = _inpand(expanded)
            clip = eroded
            clip = _kill_blobs_inner(clip, n, strength > 0)
        return clip

    return _kill_blobs_inner(clip, n, strength > 0)


def _kill_blobs_inner(clip: vs.VideoNode,
                      n: int,
                      positive: bool) -> vs.VideoNode:
    eroded = clip
    for _ in range(n):
        eroded = _inpand(eroded)

    if positive:
        result = _inflate(eroded)
        # keep only pixels present in the shrunken+expanded version
        result = _expr([clip, result], "x y min")
    else:
        result = _hysteresis(clip, eroded)

    return result


# ---------------------------------------------------------------------------
# DynTune — luma binarise mask  (DynMask=6, and as part of 2/4)
# ---------------------------------------------------------------------------

def _dyntune_mask(clip: vs.VideoNode,
                  dyn_tune: Union[int, str] = 200,
                  dyn_tune_inf: int = 0) -> vs.VideoNode:
    """
    Binarise the luma plane at *dyn_tune* threshold.

    dyn_tune can be:
      int  : simple threshold — pixels >= dyn_tune → 255
      str  : range string like "180 - 200" → keep only pixels in [low, high]

    dyn_tune_inf : inflate the resulting mask by this many pixels
                   (-1 means inpand by 1).
    """
    gray = _to_gray(clip)

    if isinstance(dyn_tune, str):
        # Range mask: "low - high"
        parts = [p.strip() for p in dyn_tune.split("-")]
        lo, hi = int(parts[0]), int(parts[1])
        mask = _expr(gray, f"x {lo} >= x {hi} <= and 255 0 ?")
    else:
        mask = _binarise(gray, int(dyn_tune))

    if dyn_tune_inf > 0:
        mask = _inflate(mask, dyn_tune_inf)
    elif dyn_tune_inf < 0:
        mask = _inpand(mask, abs(dyn_tune_inf))

    return mask


# ---------------------------------------------------------------------------
# DynColor — colour-key mask  (optional, needs vs-placebo or similar)
# ---------------------------------------------------------------------------

def _dyncolor_mask(clip: vs.VideoNode,
                   color_rgb: int,
                   tolerance: int = 10,
                   inflate: int = 1) -> vs.VideoNode:
    """
    Create a mask based on a single target colour.

    color_rgb : RGB hex value, e.g. 0xFFFFFF for white.
    tolerance : colour tolerance (1-140).
    inflate   : inflate result by this many pixels.

    Implementation uses std.Expr on an RGB24 intermediate.
    For multiple colours use the string form (not yet implemented here).
    """
    r = (color_rgb >> 16) & 0xFF
    g = (color_rgb >>  8) & 0xFF
    b = (color_rgb      ) & 0xFF

    rgb = core.resize.Bilinear(clip, format=vs.RGB24,
                               matrix_in_s="709",
                               range_in_s="limited")
    # Each plane in RGB24: plane 0=G, 1=B, 2=R  (GBR packing in VS)
    # Compute per-channel distance and threshold
    tol = tolerance
    mask = _expr(
        core.std.SplitPlanes(rgb),
        expr=[f"x {g} - abs {tol} <= y {b} - abs {tol} <= z {r} - abs {tol} <= and and 255 0 ?"]
    )
    mask = core.std.ShufflePlanes(mask, 0, vs.GRAY)
    mask = core.resize.Bilinear(mask, format=vs.GRAY8)

    if inflate > 0:
        mask = _inflate(mask, inflate)

    return mask


# ---------------------------------------------------------------------------
# DynTEdge — gradient-magnitude edge mask  (optional, needs tedgemask)
# ---------------------------------------------------------------------------

def _dyntedge_mask(clip: vs.VideoNode,
                   threshold_y: int = 20,
                   inflate: int = 2) -> vs.VideoNode:
    """
    Gradient-magnitude dynamic mask.  Requires the tedgemask plugin.

    threshold_y : magnitude threshold (2-100).
    inflate     : inflate result by this many pixels.
    """
    gray = _to_gray(clip)
    try:
        edge = core.tedgemask.TEdgeMask(gray, threshold=threshold_y, planes=[0])
    except AttributeError:
        raise RuntimeError(
            "DynTEdge path requires the tedgemask plugin.  "
            "Install it or set dyn_tedge=False."
        )
    edge = _to_gray(edge)
    edge = _binarise(edge, 1)  # anything non-zero → white

    if inflate > 0:
        edge = _inflate(edge, inflate)

    return edge


# ---------------------------------------------------------------------------
# DynMask=4 halo mask  (DynMask4H)
# ---------------------------------------------------------------------------

def _dynmask4h_mask(clip: vs.VideoNode,
                    dynmask4h: int = 60) -> vs.VideoNode:
    """
    Mask for dark subtitle outlines/halos.

    The halo is darker than the background, so we detect it by inverting:
    pixels BELOW dynmask4h threshold (i.e. dark pixels) become white in
    the mask.  This is then merged with the bright text core via hysteresis
    so only dark pixels adjacent to bright text survive.
    """
    gray = _to_gray(clip)
    # Invert: dark pixels (outline) → white, bright pixels → black
    return _expr(gray, f"x {dynmask4h} <= 255 0 ?")


# ---------------------------------------------------------------------------
# Temporal difference mask
# ---------------------------------------------------------------------------

def _temporal_diff_mask(clip: vs.VideoNode,
                        threshold: int = 10,
                        inflate: int = 1) -> vs.VideoNode:
    """
    Detect pixels that changed between adjacent frames.

    Subtitles appear and disappear; static background produces near-zero
    diff.  ANDing this with a luma mask eliminates background bleed almost
    entirely.

    threshold : minimum luma change to count as a subtitle pixel (1-255).
                Lower = more sensitive but catches more noise.
                Start around 10-20 for typical broadcast subs.
    inflate   : grow the diff mask slightly to cover full glyph bodies,
                since only the edges of appearing text change between frames.
    """
    gray = _to_gray(clip)

    # Build a 1-frame-shifted version by duplicating the first frame at the
    # start (so frame N is compared with frame N-1, no frame is dropped).
    prev = core.std.DuplicateFrames(gray, frames=[0])
    prev = core.std.Trim(prev, first=0, last=gray.num_frames - 1)

    diff = _expr([gray, prev], "x y - abs")
    mask = core.std.Binarize(diff, threshold=threshold)

    if inflate > 0:
        mask = _inflate(mask, inflate)

    return mask


# ---------------------------------------------------------------------------
# Main public function
# ---------------------------------------------------------------------------

def subtitle_mask(
    clip: vs.VideoNode,
    # Crop / location: pixels to remove from each side (left, top, right, bottom)
    loc: Optional[tuple[int, int, int, int]] = None,

    # Which mask method to use
    dyn_mask_mode: int = 2,          # 2=combo, 4=halo subs, 6=luma only

    # DynTune — luma binarise
    dyn_tune: Union[int, str] = 200,
    dyn_tune_inf: int = 0,

    # DynColor — colour key (DynMask modes 2, 5)
    dyn_color: Optional[int] = None,  # RGB hex; None = disabled
    dyn_color_tol: int = 10,
    dyn_color_inf: int = 1,

    # DynTEdge — gradient magnitude (DynMask modes 2, 7)
    dyn_tedge: bool = False,
    dyn_tedge_th_y: int = 20,
    dyn_tedge_inf: int = 2,

    # DynMask=4 halo threshold
    dynmask4h: int = 60,

    # Temporal difference AND mask — eliminates static background
    dyn_diff: bool = False,          # AND result with a temporal diff mask
    dyn_diff_threshold: int = 10,    # minimum luma change to count as subtitle
    dyn_diff_inflate: int = 1,       # inflate diff mask to cover glyph bodies

    # Global inflate after combining
    dyn_inflate: int = 0,

    # Artifact removal
    kill_noise: int = 0,
    kill_blobs: int = 0,
    pre_blobs: int = 0,
) -> vs.VideoNode:
    """
    Generate a binary subtitle mask clip.

    Parameters
    ----------
    clip : vs.VideoNode
        Input video.  Any format — internally converted to GRAY8 as needed.
    loc : (left, top, right, bottom) or None
        Pixels to remove from each side — same as VapourSynth std.Crop.
        e.g. loc=(100, 1380, 60, 500) removes 100px from left, 1380 from
        top, 60 from right, 500 from bottom.  None = full frame.
    dyn_mask_mode : int
        2  — Combo: DynTune AND optionally DynColor AND/or DynTEdge
        4  — Subtitle-with-halo: DynTune merged with DynMask4H via hysteresis
        6  — Luma binarise only (DynTune)
    dyn_tune : int or str
        Binarise threshold (1–254).  String form "lo - hi" for a luma range.
    dyn_tune_inf : int
        Inflate (+) or inpand (–1) the DynTune mask. Range: -1 to 4.
    dyn_color : int or None
        RGB hex colour for colour-key masking.  None disables this path.
        When set, this mask is logically ANDed with the DynTune mask.
    dyn_color_tol : int
        Colour tolerance (1–140).
    dyn_color_inf : int
        Inflate the DynColor mask (0–4).
    dyn_tedge : bool
        Enable gradient-magnitude edge mask (requires tedgemask plugin).
    dyn_tedge_th_y : int
        Edge magnitude threshold (2–100).
    dyn_tedge_inf : int
        Inflate the edge mask (0–8).
    dynmask4h : int
        Halo threshold for DynMask=4 (1–200).
    dyn_diff : bool
        AND the mask with a temporal difference mask.  Pixels that do not
        change between frames (static background) are removed.  This is
        the most effective way to eliminate background bleed.
        Note: causes the first frame to always be empty (no previous frame).
    dyn_diff_threshold : int
        Minimum per-pixel luma change to be considered a subtitle pixel (1–255).
        Start around 10–20; lower catches more but also more noise.
    dyn_diff_inflate : int
        Inflate the diff mask to cover full glyph bodies (only edges change).

    dyn_inflate : int
        Final inflate of the combined mask (0–10).
        Note: forced to 0 when dyn_mask_mode is 5, 6, or 7 (matches AVS behaviour).
    kill_noise : int
        Small-artifact removal.  0 = off, 1–5 = increasingly aggressive.
        Negative values pre-apply expand→inpand. See module docstring.
    kill_blobs : int
        Blob removal strength.  0 = off, positive = safer, negative = full remove.
    pre_blobs : int
        Pre-sponginess reduction for KillBlobs.  0 = off, ±1 = 'both', ±2 = 'square'.

    Returns
    -------
    vs.VideoNode
        GRAY8 binary mask (0 or 255), cropped to the loc area (or full frame).
    """

    # ---- Step 1: optional spatial crop ----
    loc_left = loc_top = loc_right = loc_bottom = 0

    if loc is not None:
        loc_left, loc_top, loc_right, loc_bottom = loc
        # loc uses the same convention as AVS Crop: each value is the number
        # of pixels to remove from that side.  std.Crop takes exactly that.
        clip = core.std.Crop(clip,
                             left=loc_left, top=loc_top,
                             right=loc_right, bottom=loc_bottom)

    # ---- Step 2: build mask per mode ----

    if dyn_mask_mode == 6:
        # Pure luma binarise
        mask = _dyntune_mask(clip, dyn_tune, dyn_tune_inf)
        dyn_inflate = 0  # matches AVS hard-set

    elif dyn_mask_mode == 5:
        # Pure colour-key
        if dyn_color is None:
            raise ValueError("dyn_mask_mode=5 requires dyn_color to be set.")
        mask = _dyncolor_mask(clip, dyn_color, dyn_color_tol, dyn_color_inf)
        dyn_inflate = 0  # matches AVS hard-set

    elif dyn_mask_mode == 7:
        # Pure edge gradient
        mask = _dyntedge_mask(clip, dyn_tedge_th_y, dyn_tedge_inf)
        dyn_inflate = 0  # matches AVS hard-set

    elif dyn_mask_mode == 2:
        # Combo: DynTune (required) + DynColor (optional AND) + DynTEdge (optional AND)
        mask = _dyntune_mask(clip, dyn_tune, dyn_tune_inf)

        if dyn_color is not None:
            color_mask = _dyncolor_mask(clip, dyn_color, dyn_color_tol, dyn_color_inf)
            mask = _expr([mask, color_mask], "x y and")

        if dyn_tedge:
            edge_mask = _dyntedge_mask(clip, dyn_tedge_th_y, dyn_tedge_inf)
            mask = _expr([mask, edge_mask], "x y and")

    elif dyn_mask_mode == 4:
        # Subtitle-with-dark-outline mask.
        # text_mask  : bright core of the subtitle (high threshold).
        # halo_mask  : dark outline pixels inverted to white (low threshold).
        # Hysteresis grows outward from the bright text seed into the dark
        # outline region — keeping only outline pixels that are directly
        # adjacent/connected to the bright text core.
        # clip_weak = candidate pixels to keep (halo — the larger region)
        # clip_strong = seed pixels that must be present (bright text core)
        text_mask = _dyntune_mask(clip, dyn_tune, dyn_tune_inf)
        halo_mask = _dynmask4h_mask(clip, dynmask4h)
        # Inflate text_mask slightly so it touches the adjacent dark outline
        text_seed = _inflate(text_mask, 1)
        # Hysteresis: keep halo pixels connected to the inflated text seed
        mask = _hysteresis(halo_mask, text_seed)
        # Merge: text core + connected outline
        mask = _expr([mask, text_mask], "x y max")

    else:
        raise ValueError(
            f"dyn_mask_mode={dyn_mask_mode} is not supported.  "
            "Use 2, 4, 5, 6, or 7."
        )

    # ---- Step 3: optional temporal diff AND ----
    if dyn_diff:
        diff_mask = _temporal_diff_mask(clip, dyn_diff_threshold, dyn_diff_inflate)
        mask = _expr([mask, diff_mask], "x y and")

    # ---- Step 4: blob removal (before KillNoise, matches AVS chain) ----
    if kill_blobs != 0:
        mask = _kill_blobs(mask, kill_blobs, pre_blobs)

    # ---- Step 5: noise/artifact removal ----
    mask = _kill_noise(mask, kill_noise)

    # ---- Step 6: final inflate ----
    if dyn_inflate > 0:
        mask = _inflate(mask, dyn_inflate)

    # ---- Step 7: pad back to original resolution ----
    # AddBorders fills the regions outside the loc area with black (0),
    # so the mask only marks pixels inside the subtitle region.
    if loc is not None:
        mask = core.std.AddBorders(mask,
                                   left=loc_left,   top=loc_top,
                                   right=loc_right, bottom=loc_bottom)

    return mask


# ---------------------------------------------------------------------------
# Convenience wrappers matching common AVS call patterns
# ---------------------------------------------------------------------------

def subtitle_mask_simple(
    clip: vs.VideoNode,
    loc: Optional[tuple[int, int, int, int]] = None,
    dyn_tune: Union[int, str] = 200,
    dyn_inflate: int = 0,
    kill_noise: int = 0,
) -> vs.VideoNode:
    """
    Simplest one-liner wrapper: luma binarise only (DynMask=6 equivalent).

    Example
    -------
    mask = subtitle_mask_simple(src, loc=(0, 900, 0, -0), dyn_tune=220)
    """
    return subtitle_mask(clip, loc=loc,
                         dyn_mask_mode=6,
                         dyn_tune=dyn_tune,
                         dyn_inflate=dyn_inflate,
                         kill_noise=kill_noise)


def subtitle_mask_halo(
    clip: vs.VideoNode,
    loc: Optional[tuple[int, int, int, int]] = None,
    dyn_tune: Union[int, str] = 200,
    dynmask4h: int = 60,
    dyn_inflate: int = 0,
    kill_noise: int = 0,
    kill_blobs: int = 0,
) -> vs.VideoNode:
    """
    Subtitle mask for subs with a dark halo (DynMask=4 equivalent).

    dyn_tune   : threshold for the bright subtitle text core.
    dynmask4h  : threshold for the surrounding dark halo (1–200).

    Example
    -------
    mask = subtitle_mask_halo(src, loc=(0, 850, 0, -0),
                              dyn_tune=200, dynmask4h=60)
    """
    return subtitle_mask(clip, loc=loc,
                         dyn_mask_mode=4,
                         dyn_tune=dyn_tune,
                         dynmask4h=dynmask4h,
                         dyn_inflate=dyn_inflate,
                         kill_noise=kill_noise,
                         kill_blobs=kill_blobs)