SRFComb2 (Vapoursynth)

"""
SRFComb2 — VapourSynth port of the AviSynth function by real.finder (Doom9).
https://forum.doom9.org/showthread.php?t=186908 + https://github.com/realfinder/AVS-Stuff/blob/master/avs%202.6%20and%20up/SRFComb.avsi
Spatial/temporal dot-crawl and rainbow suppression for analogue captures.
Handles both interlaced and progressive sources.

Requirements
------------
  mv          https://github.com/Mr-Z-2697/vapoursynth-mvtools
  zsmooth     https://github.com/adworacz/zsmooth
  warp        https://github.com/dubhater/vapoursynth-awarpsharp2
  neo_fft3d   https://github.com/HomeOfAviSynthPlusEvolution/neo_FFT3D
  fft3dfilter https://github.com/myrsloik/VapourSynth-FFT3DFilter  (fallback for neo_fft3d)
  bilateralgpu https://github.com/WolframRhodium/VapourSynth-BilateralGPU
  vszip       https://github.com/dnjulek/vapoursynth-zip
  bilateral   https://github.com/HomeOfVapourSynthEvolution/VapourSynth-Bilateral (fallback)
  llvmexpr    https://github.com/Sunflower-Dolls/Vapoursynth-llvmexpr
  akarin      https://github.com/AkarinVS                           (fallback)
  cranexpr    https://github.com/sgt0/cranexpr                      (fallback)

Optional Python modules
-----------------------
  color.py    https://github.com/Selur/VapoursynthScriptsInHybrid/blob/master/color.py
              Provides Tweak() for an accurate chroma-activity mask (ouvm).
              If not importable, a simpler approximation is used as fallback.
  sharpen.py  https://github.com/Selur/VapoursynthScriptsInHybrid/blob/master/sharpen.py
              Provides ContraSharpening() for post-degrain detail recovery.
              If not importable, contra-sharpening is silently skipped.
"""

from __future__ import annotations
import vapoursynth as vs

core = vs.core

try:
    from color import Tweak as _color_tweak  # type: ignore
except ImportError:
    _color_tweak = None

try:
    from sharpen import ContraSharpening as _contra_sharpening  # type: ignore
except ImportError:
    _contra_sharpening = None


# ---------------------------------------------------------------------------
# Plugin wrappers — each wrapper tries the fastest available backend first
# ---------------------------------------------------------------------------

def _fft3d(clip: vs.VideoNode, **kwargs) -> vs.VideoNode:
    """FFT3D — prefers neo_fft3d, falls back to fft3dfilter."""
    if hasattr(core, "neo_fft3d"):
        return core.neo_fft3d.FFT3D(clip, **kwargs)
    return core.fft3dfilter.FFT3DFilter(clip, **kwargs)


def _bilateral(clip: vs.VideoNode, sigmaS: float = 3.0, sigmaR: float = 0.02, **kwargs) -> vs.VideoNode:
    """Bilateral filter — prefers bilateralgpu, then vszip, then bilateral.
    All backends share the sigmaS / sigmaR parameter names.
    Note: there is no 'diameter' parameter; spatial radius is controlled by sigmaS.
    """
    if hasattr(core, "bilateralgpu"):
        return core.bilateralgpu.Bilateral(clip, sigmaS=sigmaS, sigmaR=sigmaR, **kwargs)
    if hasattr(core, "vszip"):
        return core.vszip.Bilateral(clip, sigmaS=sigmaS, sigmaR=sigmaR, **kwargs)
    if hasattr(core, "bilateral"):
        return core.bilateral.Bilateral(clip, sigmaS=sigmaS, sigmaR=sigmaR, **kwargs)
    raise RuntimeError(
        "srfcomb2: a bilateral filter plugin is required (bilateralgpu, vszip, or bilateral) — "
        "install one from https://github.com/dnjulek/vapoursynth-zip"
    )


def _box_blur(
    clip: vs.VideoNode,
    hradius: int = 1,
    hpasses: int = 1,
    vradius: int = 1,
    vpasses: int = 1,
    planes: list[int] | None = None,
) -> vs.VideoNode:
    """BoxBlur — prefers vszip, falls back to std."""
    kwargs: dict = dict(hradius=hradius, hpasses=hpasses, vradius=vradius, vpasses=vpasses)
    if planes is not None:
        kwargs["planes"] = planes
    if hasattr(core, "vszip"):
        return core.vszip.BoxBlur(clip, **kwargs)
    return core.std.BoxBlur(clip, **kwargs)


def _temporal_soften(
    clip: vs.VideoNode,
    radius: int,
    luma_threshold: int,
    chroma_threshold: int,
    scenechange: int,
    mode: int,
) -> vs.VideoNode:
    """TemporalSoften — zsmooth or focus2 as fallback.

    zsmooth signature:  (clip, radius, threshold[], scenechange)
    focus2 signature:   (clip, radius, luma_threshold, chroma_threshold, scenechange, mode)
    """
    if hasattr(core, "focus2"):
        return core.focus2.TemporalSoften2(
            clip, radius, luma_threshold, chroma_threshold, scenechange, mode
        )
    # zsmooth uses a single threshold array [luma, chroma]; no mode parameter
    return core.zsmooth.TemporalSoften(
        clip, radius, [luma_threshold, chroma_threshold], scenechange
    )


def _expr(clips: vs.VideoNode | list[vs.VideoNode], expr: str | list[str]) -> vs.VideoNode:
    """Expr — prefers llvmexpr → akarin → cranexpr → std."""
    if hasattr(core, "llvmexpr"):
        return core.llvmexpr.Expr(clips, expr)
    if hasattr(core, "akarin"):
        return core.akarin.Expr(clips, expr)
    if hasattr(core, "cranexpr"):
        return core.cranexpr.Expr(clips, expr)
    return core.std.Expr(clips, expr)


def _checkmate(clip: vs.VideoNode) -> vs.VideoNode:
    """Checkmate deinterlacing-artefact removal via vszip (no-op if unavailable)."""
    if not hasattr(core, "vszip"):
        return clip
    bits = clip.format.bits_per_sample
    if bits != 8:
        fmt8 = clip.format.replace(bits_per_sample=8)
        clip8 = core.resize.Point(clip, format=fmt8)
        result = core.vszip.Checkmate(clip8)
        return core.resize.Point(result, format=clip.format)
    return core.vszip.Checkmate(clip)


# ---------------------------------------------------------------------------
# Internal helpers (mt_* equivalents)
# ---------------------------------------------------------------------------

def _mt_logic(a: vs.VideoNode, b: vs.VideoNode, mode: str = "min") -> vs.VideoNode:
    op = "x y min" if mode == "min" else "x y max"
    return _expr([a, b], op)


def _mt_binarize(clip: vs.VideoNode, threshold: int) -> vs.VideoNode:
    peak = (1 << clip.format.bits_per_sample) - 1
    return _expr([clip], f"x {threshold} > {peak} 0 ?")


def _mt_expand(clip: vs.VideoNode, n: int = 1) -> vs.VideoNode:
    for _ in range(n):
        clip = core.std.Maximum(clip)
    return clip


def _mt_expand_horizontal(clip: vs.VideoNode, n: int = 1) -> vs.VideoNode:
    for _ in range(n):
        clip = core.std.Maximum(clip, coordinates=[0, 0, 0, 1, 1, 0, 0, 0])
    return clip


def _mt_inflate(clip: vs.VideoNode, n: int = 1) -> vs.VideoNode:
    for _ in range(n):
        clip = core.std.Inflate(clip)
    return clip


def _luma(clip: vs.VideoNode) -> vs.VideoNode:
    return core.std.ShufflePlanes(clip, 0, vs.GRAY)


def _weave_fields(clip: vs.VideoNode) -> vs.VideoNode:
    """DoubleWeave then drop odd frames — equivalent to AviSynth Weave()."""
    return core.std.DoubleWeave(clip)[::2]


def _scale_chroma_mask(
    mask: vs.VideoNode,
    target_width: int,
    target_height: int,
    is_420_or_422: bool,
) -> vs.VideoNode:
    """Resize a chroma mask with the correct sub-sampling offset."""
    src_left = -0.25 if is_420_or_422 else 0.0
    return core.resize.Bilinear(mask, target_width, target_height, src_left=src_left)


# ---------------------------------------------------------------------------
# Public function
# ---------------------------------------------------------------------------

def srfcomb2(
    clip: vs.VideoNode,
    *,
    dedot: bool | None = None,
    derainbow: bool | None = None,
    thsad: int | None = None,
    thsadc: int | None = None,
    pal: bool | None = None,
    progressive: bool | None = None,
    contrasharp: bool = True,
) -> vs.VideoNode:
    """
    SRFComb2 — spatial + temporal dot-crawl and rainbow artefact reduction.

    Ported from the AviSynth function by real.finder (Doom9), v1.01.

    Parameters
    ----------
    clip : vs.VideoNode
        YUV planar input. Both interlaced and progressive are supported.
        For interlaced material the clip should *not* be pre-separated;
        fields are separated internally.
    dedot : bool, optional
        Enable luma (dot-crawl) reduction. Defaults to *progressive*.
    derainbow : bool, optional
        Enable chroma (rainbow) reduction. Defaults to *progressive*.
    thsad : int, optional
        SAD threshold for dot-crawl motion compensation.
        Defaults to 700 (progressive) or 500 (interlaced).
    thsadc : int, optional
        SAD threshold for rainbow motion compensation.
        Defaults to 700 (progressive) or 500 (interlaced).
    pal : bool, optional
        Hint that the source is PAL (25 fps). Affects mask expansion
        direction. Auto-detected from frame-rate when not given.
    progressive : bool, optional
        Force progressive (``True``) or interlaced (``False``) handling.
        When omitted, auto-detected from the ``_FieldBased`` frame
        property (0 = progressive, 1 = BFF, 2 = TFF). Override this
        if the source is mistagged.
    contrasharp : bool
        Apply contra-sharpening after motion compensation to recover
        detail lost during degrain (default ``True``).

    Returns
    -------
    vs.VideoNode
        Processed clip in the same format as the input.

    Raises
    ------
    ValueError
        If the input is not planar YUV, or is greyscale.
    """

    if clip.format.color_family != vs.YUV:
        raise ValueError("srfcomb2: YUV planar input required")
    if clip.format.num_planes == 1:
        raise ValueError("srfcomb2: greyscale input is not supported")

    # ------------------------------------------------------------------
    # Resolve defaults
    # ------------------------------------------------------------------
    if progressive is None:
        field_based = clip.get_frame(0).props.get('_FieldBased', 0)
        progressive = field_based == 0
    else:
        field_based = 0 if progressive else 2  # assume TFF if forced interlaced

    dedot     = progressive if dedot     is None else dedot
    derainbow = progressive if derainbow is None else derainbow
    thsad     = (700 if progressive else 500) if thsad  is None else thsad
    thsadc    = (700 if progressive else 500) if thsadc is None else thsadc
    if pal is None:
        fps = clip.fps_num / clip.fps_den if clip.fps_den else 0
        pal = abs(fps - 25.0) < 0.1

    ss     = clip.format.subsampling_w  # 0 = 4:4:4, 1 = 4:2:2 / 4:2:0
    is_sub = ss > 0                     # True for 420 / 422

    # ------------------------------------------------------------------
    # Field separation (interlaced path)
    # ------------------------------------------------------------------
    fields = clip if progressive else core.std.SeparateFields(clip, tff=field_based == 2)
    oY = _luma(fields)

    # ------------------------------------------------------------------
    # Luma edge / comb mask
    # ------------------------------------------------------------------
    e_h = core.std.Convolution(oY, [0, 0, 0, -16, 0, 16, 0, 0, 0], saturate=False)
    e_v = core.std.Convolution(oY, [0,  0,  0,  0, 2, -2, 0, 0, 0], saturate=False)

    luma_edge = _mt_logic(e_h, e_v, "min")
    luma_mask = _mt_logic(luma_edge, core.std.Invert(oY), "min")

    # AviSynth uses horizontal-only expand for PAL; vertical+horizontal for NTSC
    if pal:
        luma_mask = _mt_expand(luma_mask, 2)
    else:
        luma_mask = _mt_expand_horizontal(luma_mask, 2)
    luma_mask = _mt_inflate(luma_mask)

    # ------------------------------------------------------------------
    # UV difference mask (ouvm) — chroma activity map
    #
    # AviSynth: MakeDiff(Tweak(sat=0), Tweak(sat=PAL?20:10))
    # The diff between a fully desaturated clip and a heavily boosted one
    # is zero where chroma is already neutral and grows with colour content.
    # The U and V channels of that diff are then mapped through:
    #   x == range_half → 0,  else abs(range_half - x) * 2.28
    # and the per-pixel max of U and V is used as the activity signal.
    # ------------------------------------------------------------------
    mid = 1 << (clip.format.bits_per_sample - 1)

    if _color_tweak is not None:
        # Accurate path — uses color.Tweak
        tweak_zero  = _color_tweak(fields, sat=0)
        tweak_boost = _color_tweak(fields, sat=20 if pal else 10)
        chroma_diff = core.std.MakeDiff(tweak_zero, tweak_boost)
        u_diff_raw  = core.std.ShufflePlanes(chroma_diff, 1, vs.GRAY)
        v_diff_raw  = core.std.ShufflePlanes(chroma_diff, 2, vs.GRAY)
        uv_diff     = core.std.Interleave([u_diff_raw, v_diff_raw])
        uv_mapped   = _expr([uv_diff], f"x {mid} = 0 {mid} x - abs 2.28 * ?")
        ouvm_u      = core.std.SelectEvery(uv_mapped, 2, [0])
        ouvm_v      = core.std.SelectEvery(uv_mapped, 2, [1])
        ouvm_sub    = _mt_logic(ouvm_u, ouvm_v, "max")
    else:
        # Fallback — raw plane deviation from mid-grey
        u_plane  = core.std.ShufflePlanes(fields, 1, vs.GRAY)
        v_plane  = core.std.ShufflePlanes(fields, 2, vs.GRAY)
        u_dev    = _expr([u_plane], f"x {mid} - abs 2.28 *")
        v_dev    = _expr([v_plane], f"x {mid} - abs 2.28 *")
        ouvm_sub = _mt_logic(u_dev, v_dev, "max")

    # Resize ouvm to luma dimensions for mask combination
    ouvm      = _scale_chroma_mask(ouvm_sub, oY.width, oY.height, is_sub)
    luma_mask = _mt_logic(luma_mask, _mt_binarize(ouvm, 30), "min")

    # ------------------------------------------------------------------
    # Luma denoise (FFT3D)
    # ------------------------------------------------------------------
    cle_y = _fft3d(
        oY,
        sigma =15 if progressive else 33,
        sigma2=15 if progressive else 33,
        sigma3=10 if progressive else 22,
        sigma4=0,
        bt=1,
    )

    if progressive:
        blur   = _box_blur(oY, hradius=1, vradius=1)
        detail = core.std.MakeDiff(oY, blur)
        cle_y  = core.std.MergeDiff(cle_y, _expr([detail], "x 0.25 *"))
    else:
        # aBlur(blur=2) ≈ horizontal-only blur
        cle_y = core.std.Merge(cle_y, core.warp.ABlur(oY, blur=2))

    # Weave back before repair (interlaced path)
    cle_y_woven = cle_y if progressive else _weave_fields(cle_y)

    # AviSynth uses Checkmate → TemporalRepair → Repair
    cle_y_woven = _checkmate(cle_y_woven)
    cle_y_woven = core.zsmooth.TemporalRepair(cle_y_woven, cle_y_woven, mode=3)
    if not progressive:
        cle_y_sep = core.std.SeparateFields(cle_y_woven)
        cle_y_woven = core.zsmooth.Repair(cle_y_sep, cle_y, mode=3)
        cle_y_woven = _weave_fields(cle_y_woven)
    else:
        cle_y_woven = core.zsmooth.Repair(cle_y_woven, oY, mode=3)

    # Replace luma in fields clip
    luma_cleaned_yuv = core.std.ShufflePlanes(
        [cle_y_woven, fields, fields], [0, 1, 2], vs.YUV
    )
    merged = core.std.MaskedMerge(fields, luma_cleaned_yuv, luma_mask)

    # ------------------------------------------------------------------
    # Chroma comb mask
    # ------------------------------------------------------------------
    chm_h = core.std.Convolution(oY, [0, 0, 0, -1, 0, 1, 0, 0, 0], saturate=False)
    chm_v = core.std.Convolution(oY, [0,  0,  0,  0, 2, -2, 0, 0, 0], saturate=False)

    chm = _mt_logic(chm_h, chm_v, "min")
    chm = _mt_logic(chm, _mt_binarize(core.std.Invert(oY), 50), "min")
    chm = _mt_expand_horizontal(chm)
    chm = _mt_inflate(chm)

    # ------------------------------------------------------------------
    # Chroma denoise (FFT3D, sigma2/3/4 only → chroma planes)
    # ------------------------------------------------------------------
    chroma_clean = _fft3d(
        merged,
        sigma =0,
        sigma2=11 if progressive else 22,
        sigma3= 5 if progressive else 11,
        sigma4=22 if progressive else 44,
        bt=1,
    )
    spati_comb_c = core.std.MaskedMerge(merged, chroma_clean, chm)

    # ------------------------------------------------------------------
    # UV bilateral smoothing
    # ------------------------------------------------------------------
    u_in = core.std.ShufflePlanes(spati_comb_c, 1, vs.GRAY)
    v_in = core.std.ShufflePlanes(spati_comb_c, 2, vs.GRAY)
    uv_interleaved = core.std.Interleave([u_in, v_in])
    uv_filtered    = _bilateral(uv_interleaved, sigmaS=1.4, sigmaR=0.028)
    u_filtered = core.std.SelectEvery(uv_filtered, 2, [0])
    v_filtered = core.std.SelectEvery(uv_filtered, 2, [1])
    spati_comb_c = core.std.ShufflePlanes(
        [spati_comb_c, u_filtered, v_filtered], [0, 0, 0], vs.YUV
    )

    # ------------------------------------------------------------------
    # aWarpSharp2 chroma sharpening
    # The AviSynth original used chroma=6 (process chroma only, copy luma)
    # with separate H/V depth passes. The VS plugin only accepts chroma=0
    # or chroma=1, but supports a planes parameter to restrict processing.
    # chroma=1 uses independent per-channel edge masks (closest to AviSynth
    # chroma=3/6). planes=[1,2] processes only U and V, copying luma.
    # ------------------------------------------------------------------
    spati_comb_c = core.warp.AWarpSharp2(spati_comb_c, depth=[16, 8, 8], chroma=1, planes=[1, 2])

    # Combined luma+chroma combmask for the final merge
    chm_uv_resized = _scale_chroma_mask(chm, ouvm_sub.width, ouvm_sub.height, is_sub)
    full_comb_mask = core.std.ShufflePlanes(
        [luma_mask, chm_uv_resized, chm_uv_resized], [0, 0, 0], vs.YUV
    )

    # ------------------------------------------------------------------
    # Motion compensation — separate search super and degrain super
    # (AviSynth uses rfilter=4 for search, levels=1 for degrain)
    # ------------------------------------------------------------------
    if progressive:
        super_search  = core.mv.Super(spati_comb_c, rfilter=4)
        super_degrain = core.mv.Super(fields, levels=1)

        bv = core.mv.Analyse(super_search, isb=True,  delta=1, overlap=2, dct=8)
        fv = core.mv.Analyse(super_search, isb=False, delta=1, overlap=2, dct=8)

        degrain = core.mv.Degrain1(fields, super_degrain, bv, fv,
                                   thsad=thsad, thsadc=thsadc)
    else:
        even_sc = core.std.SelectEvery(spati_comb_c, 2, [0])
        odd_sc  = core.std.SelectEvery(spati_comb_c, 2, [1])
        even_o  = core.std.SelectEvery(fields, 2, [0])
        odd_o   = core.std.SelectEvery(fields, 2, [1])

        sup_se = core.mv.Super(even_sc, rfilter=4)
        sup_so = core.mv.Super(odd_sc,  rfilter=4)
        sup_de = core.mv.Super(even_o,  levels=1)
        sup_do = core.mv.Super(odd_o,   levels=1)

        bv_e = core.mv.Analyse(sup_se, isb=True,  delta=1, overlap=2, dct=8)
        fv_e = core.mv.Analyse(sup_se, isb=False, delta=1, overlap=2, dct=8)
        bv_o = core.mv.Analyse(sup_so, isb=True,  delta=1, overlap=2, dct=8)
        fv_o = core.mv.Analyse(sup_so, isb=False, delta=1, overlap=2, dct=8)

        even_d = core.mv.Degrain1(even_o, sup_de, bv_e, fv_e, thsad=thsad, thsadc=thsadc)
        odd_d  = core.mv.Degrain1(odd_o,  sup_do, bv_o, fv_o, thsad=thsad, thsadc=thsadc)
        degrain = core.std.Interleave([even_d, odd_d])

    # ------------------------------------------------------------------
    # Motion mask — only when dedot or derainbow is active
    # AviSynth combines bv AND fv masks; old port only used bv
    # ------------------------------------------------------------------
    if dedot or derainbow:
        if progressive:
            o_for_mask = fields
            mmask_bv = core.mv.Mask(o_for_mask, bv, kind=0, ml=1, ysc=255)
            mmask_fv = core.mv.Mask(o_for_mask, fv, kind=0, ml=1, ysc=255)
        else:
            mmask_bv_e = core.mv.Mask(even_o, bv_e, kind=0, ml=1, ysc=255)
            mmask_fv_e = core.mv.Mask(even_o, fv_e, kind=0, ml=1, ysc=255)
            mmask_bv_o = core.mv.Mask(odd_o,  bv_o, kind=0, ml=1, ysc=255)
            mmask_fv_o = core.mv.Mask(odd_o,  fv_o, kind=0, ml=1, ysc=255)
            mmask_e    = _mt_logic(mmask_bv_e, mmask_fv_e, "max")
            mmask_o    = _mt_logic(mmask_bv_o, mmask_fv_o, "max")
            mmask_bv   = core.std.Interleave([mmask_e, mmask_o])
            mmask_fv   = mmask_bv  # already combined above

        motion_mask = _mt_logic(mmask_bv, mmask_fv, "max") if progressive else mmask_bv

        # TemporalSoften + expand
        ts = _temporal_soften(motion_mask, 2, 255, 255, 0, 2)
        motion_mask = _mt_logic(motion_mask, _mt_expand(ts, 1), "max")
        motion_mask = _mt_expand(motion_mask, 1)
        motion_mask = _mt_expand_horizontal(motion_mask)
        motion_mask = _mt_inflate(motion_mask)

        if contrasharp and _contra_sharpening is not None:
            sharp = _contra_sharpening(spati_comb_c, degrain)
        else:
            sharp = spati_comb_c

        # Selective plane merging: dedot controls Y, derainbow controls U/V
        y_plane  = 3 if dedot      else 2
        uv_plane = 3 if derainbow  else 2
        merged = core.std.MaskedMerge(
            degrain, sharp, motion_mask,
            planes=[0] if dedot and not derainbow else
                   [1, 2] if derainbow and not dedot else
                   [0, 1, 2],
        )
    else:
        merged = degrain

    # ------------------------------------------------------------------
    # Final merge: blend result back over original via full combmask
    # ------------------------------------------------------------------
    final = core.std.MaskedMerge(fields, merged, full_comb_mask)

    return final if progressive else _weave_fields(final)