import vapoursynth as vs
from fvsfunc import OverlayInter, Depth
from nnedi3_rpow2 import *
from vardefunc import deband
from lvsfunc import mask as dtlmask
from hysteria import Hysteria
from mvsfunc import BM3D
from havsfunc import QTGMC
from dhh import mask as dehalo_mask
from fine_dehalo import *
from xaa import *
from kagefunc import adaptive_grain
from vsutil import join, plane
from mvsfunc import BM3D
core = vs.core
core.max_cache_size = 14000

#Source
src = core.d2v.Source(r"G:\Encode\.vscode\Scripts\[BD-DVD]\Romeo x Juliet\src\01-02\VTS_01_1.d2v")[42558:85986]
video = src
video = Depth(video, 16)
ep = video[0:38975]
ed = video[38975:41670]
av = video[41670:43428]

deband_args = dict(iterations=2, threshold=2.0, radius=13)

#IVTC
IVTC_ep = core.tivtc.TFM(ep, mode=6, PP=4, field=1, slow=2, cthresh=8, micmatching=3, clip2=QTGMC(ep, Preset='Very Slow', SourceMatch=2, TFF=True, FPSDivisor=2, Sharpness=0.4, TR2=3))
IVTC_ep = core.tivtc.TDecimate(IVTC_ep, mode=1)
IVTC_ed = OverlayInter(ed, pattern=0, tff=True)
IVTC_av = core.tivtc.TFM(av, mode=6, PP=4, field=1, slow=2, cthresh=8, micmatching=3, clip2=QTGMC(av, Preset='Very Slow', SourceMatch=2, TFF=True, FPSDivisor=2, Sharpness=0.4, TR2=3))
IVTC_av = core.tivtc.TDecimate(IVTC_av, mode=1)
IVTC = IVTC_ep + IVTC_ed + IVTC_av

#Crop
crop = core.std.Crop(IVTC, top=2, left=4, right=2)
crop = core.edgefixer.Continuity(crop, right=2, top=2, radius=6)
crop = nnedi3_rpow2(crop, pscrn=1, width=848, height=480, nns=4, qual=2)
mask = Hysteria(crop)
crop = core.std.MaskedMerge(crop, mask, mask)

#AA
aa_a = xaa(crop, mode="di3 nnedi3cl", mask=1, mthr=8, nns=4, uscl="nnedi3cl", dscl="Spline36")
aa = core.std.MaskedMerge(aa_a, crop, crop)

#Dehalo
dh_mask = dehalo_mask(aa, smooth=True, dha=False)
dehalo = fine_dehalo(aa, rx=2.0, darkstr=0)
masked_dehalo = core.std.MaskedMerge(dehalo, aa, dh_mask)

#Denoise
denoise = BM3D(masked_dehalo, sigma=[3,3,3], radius1=1, profile1='fast')
denoise_mask = dtlmask.detail_mask(masked_dehalo, rad=1, brz_a=0.029, brz_b=0.057)
denoised = core.std.MaskedMerge(denoise, masked_dehalo, denoise_mask)

#Deband
deband = join([
    core.placebo.Deband(plane(denoised, 0), grain=4, **deband_args),
    core.placebo.Deband(plane(denoised, 1), grain=0, **deband_args),
    core.placebo.Deband(plane(denoised, 2), grain=0, **deband_args)
])
debanded = core.std.MaskedMerge(deband, denoised, denoise_mask)

#Grain
grain = adaptive_grain(debanded, 0.25)

#Final
out = grain
final = Depth(out, 16)
final.set_output()