from vssource import BestSource
from vstools import initialize_clip, Matrix, Primaries, Transfer
from vsview import set_output, is_preview
from vskernels import EwaLanczosSharp
from vsmasktools import HardsubSign, BoundingBox

bili_4k = BestSource.source("[Doomdos] - 记忆管理局 - 第1话 - [4K].mp4")
bili_soft = BestSource.source("FALSE.MEMORY.S01E01.REPACK.1080p.BILI.WEB-DL.ZHO.AAC2.0.H.264.MSubs-ToonsHub.mkv")

bili_4k = initialize_clip(bili_4k, matrix=Matrix.BT709, primaries=Primaries.BT709, transfer=Transfer.BT709)
bili_soft = initialize_clip(bili_soft, matrix=Matrix.BT709, primaries=Primaries.BT709, transfer=Transfer.BT709)

# Cut Bilibili intros
bili_4k = bili_4k[408:]
bili_soft = bili_soft[264:41776]

upscale = EwaLanczosSharp().scale(bili_soft, 3840, 2160, sigmoid=True, antiring=0.6)

# Splice missing frame
bili_4k = bili_4k[:39966] + upscale[39966:39968] + bili_4k[39967:]

# Bilibili logo
logo = HardsubSign(bound=BoundingBox((327, 116), (3358, 114)), thr=0.03)

delogo = logo.apply_dehardsub(bili_4k, upscale)
delogo_mask = logo.get_mask(bili_4k, upscale)

# First frame containing subs, Idk of a fast way to do this
top_hardsubs = HardsubSign(ranges=[(240, 261)], bound=BoundingBox((299, 179), (1772, 40)), thr=0.1)

dehardsub = top_hardsubs.apply_dehardsub(delogo, upscale)
dehardsub_mask = top_hardsubs.get_mask(delogo, upscale)

if is_preview():
    set_output(bili_4k, "4K Source")
    set_output(bili_soft, "No Hardsubs")
    set_output(upscale, "Upscaled")
    set_output(initialize_clip(delogo_mask), "BiliBili Mask")
    set_output(delogo, "Delogoed")
    set_output(initialize_clip(dehardsub_mask), "Hardsub Mask")
    set_output(dehardsub, "Dehardsubbed")