import vapoursynth as vs
import vodesfunc.scale
import vsmasktools
from vskernels import Catrom
from vodesfunc import out
import vsmlrt
import vodesfunc as vof
from vstools import get_y, plane, depth, join
from vsdenoise import mvtools, MVToolsPresets
from rekt import rektlvls
from awsmfunc import bbmod
import vsdenoise
from adptvgrnMod import adptvgrnMod as graining
from vsrgtools import contrasharpening


core = vs.core


sp = core.lsmas.LWLibavSource(r'[SubsPlease] Kono Subarashii Sekai ni Bakuen wo! - 03 (1080p) [687845A2].mkv')
sp = depth(sp, 16)

rekt = rektlvls(sp, [0,-1], [10,10], [0,-1], [10,10])
bbm = bbmod(rekt, 1, 1, 1, 1,)

rescale = vof.vodes_rescale(bbm, 837, Catrom, doubler=vof.scale.NNEDI_Doubler(False), credit_mask=True, line_mask=True, base_width=1448, base_height=837, use_baseheight_mask=True, )[0]



sloc = [0, 0, 0.35, 1, 0.4, 2, 0.8, 32, 1, 64]
dft = core.dfttest.DFTTest(rescale, sigma=8, sigma2=16, tbsize=1, ssx=sloc, ssy=sloc, planes=[0])

mv = mvtools.MVTools(dft, tr=2, refine=2, **MVToolsPresets.SMDE())
ref = mv.degrain(thSAD=65, ref=None)
ref = depth(ref, 32)
bm3d = core.bm3dcuda_rtc.BM3Dv2(get_y(depth(dft, 32)), get_y(ref), sigma=0.5, radius=3).resize.Bicubic(format=vs.YUV420P16)

knlm = vsdenoise.knl_means_cl(dft, .6, )

#joined = join([plane(bm3d, 0), plane(knlm,1), plane(knlm,2),])

joined = core.std.ShufflePlanes([bm3d, knlm], [0,1,2], vs.YUV420P16)

ccd = vsdenoise.ccd(joined, 4)

rgb = core.resize.Bicubic(ccd, format=vs.RGBS, matrix_in=1, transfer_in=1, primaries_in=1,
                          range_in=0).std.Limiter()
dpir = vsmlrt.DPIR(rgb, 10,
                   backend=vsmlrt.Backend().TRT(fp16=True, workspace=10000, use_cuda_graph=True, use_cublas=True,
                                                num_streams=2), model=3, ).resize.Bicubic(format=vs.YUV420P16, matrix=1,
                                                                                          transfer=1, primaries=1,
                                                                                          range=0)

dpir = contrasharpening(dpir, dft, )

adg = core.tcanny.TCanny(get_y(dft), 0, 0, mode=1, op=3, scale=.95).std.Binarize(8000)
dpir = core.std.MaskedMerge(dpir, dft, adg, [0])

deband = core.neo_f3kdb.Deband(dpir, 10, 32, 32, 32, 0, 0)
deband = core.neo_f3kdb.Deband(deband, 20, 48, 40, 40, 0, 0)
debandmask = vsmasktools.detail_mask(plane(dpir, 0), .45, .15 )# .edgemask(plane(nuke, 0), lthr=80 << 8, hthr=150 << 8)
debanded = core.std.MaskedMerge(deband, dpir, debandmask)


grain = graining(debanded, [0.2, 0.8], 1.3, 75, temporal_average=15, temporal_radius=5, luma_scaling=7, grainer=lambda x: core.noise.Add(x, 2, .25, 3, xsize=2.5, ysize=2.5, constant=False))
final = depth(grain, 10)