import vapoursynth as vs
from fvsfunc import Depth, Despline36M
from lvsfunc import mask as lvfmask
from keyframes import generate_qpfile_double as qpfile
from nnedi3_rpow2 import nnedi3_rpow2
from vsutil import get_y
from mvsfunc import BM3D
from muvsfunc import MergeChroma
from keyframes import generate_qpfile_double as qp
from project_module import flt
from vardefunc import noise
from vsdehalo import fine_dehalo, edge_cleaner
core = vs.core
core.max_cache_size = 14000

#source
src = core.lsmas.LWLibavSource(r"G:\Encode\.vscode\Scripts\[BD-DVD]\Sword Art Online\src\01.m2ts")
videod = src
videod = Depth(videod, 16)
videod = videod.edgefixer.Continuity(left=2, top=2, right=2, bottom=2)

#Dehalo
dehalo = fine_dehalo(videod, rx=1.8, darkstr=0)
videod = edge_cleaner(dehalo, 5.75, hot=True, smode=True)

#Descale
descale = Despline36M(videod, 1280, 720)
video = nnedi3_rpow2(descale, nns=4, qual=2, pscrn=1).resize.Spline64(1920, 1080)

#Denoise
d_mask = lvfmask.detail_mask(videod, sigma=4, rad=2, brz_a=0.026, brz_b=0.0023)
luma = get_y(video)
denoise_Y_a = BM3D(luma, radius1=1, profile1='fast', sigma=[2.5,0,0])
denoise_Y_b = BM3D(luma, radius1=1, profile1='fast', sigma=[3.5,0,0])
denoise_a = MergeChroma(denoise_Y_a, video)
denoise_b = MergeChroma(denoise_Y_b, video)
denoised = core.std.MaskedMerge(denoise_b, denoise_a, d_mask)

#Deband
debanded = core.average.Mean([
        flt.masked_f3kdb(denoised, rad=17, thr=[20, 24], grain=[24, 12]),
        flt.masked_f3kdb(denoised, rad=21, thr=[32, 24], grain=[24, 12]),
        flt.masked_placebo(denoised, rad=6, thr=2.8, itr=2, grain=6)
    ])

#Resize
resize = core.resize.Spline36(debanded, 1280, 720).std.Merge(descale, [0,1])
clipmax = core.std.Maximum(descale)
clipmin = core.std.Minimum(descale)
f_mask = core.std.Expr([clipmax, clipmin], 'x y -').std.Binarize(2900).resize.Spline16(format=vs.GRAY16)

#Grain
detail = core.std.MaskedMerge(resize, descale, f_mask)
grain = noise.Graigasm(
    thrs=[x << 8 for x in (32, 80, 128, 176)],
    strengths=[(0.25, 0.0), (0.20, 0.0), (0.15, 0.0), (0.0, 0.0)],
    sizes=(1.25, 1.20, 1.15, 1),
    sharps=(80, 70, 60, 50),
    grainers=[
        noise.AddGrain(seed=69420, constant=False),
        noise.AddGrain(seed=69420, constant=False),
        noise.AddGrain(seed=69420, constant=True)
        ]).graining(detail)

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

qp(src, out_path=r"G:\Encode\keyframes\sao1-01.txt", autismo=4)