from vsaa import NNEDI3, based_aa, pre_aa
from vstools import ChromaLocation, ColorRange, Matrix, Primaries, Transfer, vs

from soifunc.resize import good_resize

__all__ = [
    "based_imagesub",
]


def based_imagesub(clip: vs.VideoNode, sub_file: str) -> vs.VideoNode:
    """
    ImageSub that performs anti-aliasing on only the subtitles.

    :param clip: Video to add subtitles to
    :type clip: vs.VideoNode
    :param sub_file: Path to the image-based subtitles file
    :type sub_file: str
    :return: Subtitled video
    :rtype: vs.VideoNode
    """
    matrix = Matrix.from_video(clip)
    primaries = Primaries.from_video(clip)
    transfer = Transfer.from_video(clip)
    range = ColorRange.from_video(clip)
    sub = clip.sub.ImageFile(
        file=sub_file,
        blend=False,
    )
    sub = sub.resize.Spline36(
        format=clip.format.id,
        matrix=matrix,
        primaries_in=primaries,
        primaries=primaries,
        transfer_in=transfer,
        transfer=transfer,
        range=range,
    )
    sub = good_resize(sub, clip.width, clip.height)
    # Use NNEDI3 since ArtCNN isn't really designed for text
    sub = based_aa(sub, prefilter=pre_aa, supersampler=NNEDI3)

    mask = sub.std.PropToClip(prop="_Alpha")
    mask_format = vs.core.query_video_format(
        mask.format.color_family,
        clip.format.sample_type,
        clip.format.bits_per_sample,
        mask.format.subsampling_w,
        mask.format.subsampling_h,
    )
    mask = mask.resize.Spline36(
        clip.width,
        clip.height,
        format=mask_format.id,
    )

    return vs.core.std.MaskedMerge(clipa=clip, clipb=sub, mask=mask)
