import vapoursynth as vs
core = vs.core
import lvsfunc as lvf
import kagefunc as kgf
import havsfunc as hvf
import vardefunc as vdf
import debandshit as dbs
import EoEfunc as eoe
from vsutil import get_y, get_w, split, join
from adptvgrnMod import adptvgrnMod
from jvsfunc import ccd
#from ccd import ccd
###CPU management
import psutil
p_handle = psutil.Process()
#p_handle.cpu_affinity([0, 2, 4, 6, 8, 10, 12, 14]) #p-cores no mt
#p_handle.cpu_affinity([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]) #p-cores mt
#p_handle.cpu_affinity([0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 19]) #p+e cores no mt minus 1 e-core
#p_handle.cpu_affinity([0, 2, 4, 6, 8, 10, 12, 14, 16, 17, 18, 19]) #p+e cores no mt
p_handle.cpu_affinity([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19]) #p+e cores mt
core.num_threads = 20
###

###Ranges
skip_aa = [(0, 520), (133925, 139565)]
###

###Source
src   = core.lsmas.LWLibavSource(r'E:\Entertainment\Anime\SAO-P\BD_VIDEO\BDMV\STREAM\00002.m2ts').std.SetFrameProp(prop='_Matrix', intval=1).std.SetFrameProp(prop='_Transfer', intval=1).std.SetFrameProp(prop='_ColorRange', intval=1).std.SetFrameProp(prop='_Primaries', intval=1)
encode   = core.lsmas.LWLibavSource(r'E:\Entertainment\Anime\SAO-P\BD_VIDEO\BDMV\STREAM\SAO-P.hevc')
src32 = core.fmtc.bitdepth(src, bits=32)
###

###Anti-Aliasing
y, u, v, = split(src32)
#Supersample AA
nn3ss_v = core.nnedi3cl.NNEDI3CL(y, 1, dh=True, dw=False, nsize=5, nns=4, qual=2, pscrn=1)
nn3ss_h = core.nnedi3cl.NNEDI3CL(nn3ss_v.std.Transpose(), 1, dh=True, dw=False, nsize=5, nns=4, qual=2, pscrn=1)
nn3ssaa = core.std.Transpose(nn3ss_h).resize.Bicubic(src_top=-.5, src_left=-.5)
ee3ss_v = core.eedi3m.EEDI3CL(y, 1, dh=True, alpha=0.6, beta=0.2, gamma=50, nrad=2, mdis=12, sclip=nn3ss_v).std.Transpose()
ee3ss_h = core.eedi3m.EEDI3CL(ee3ss_v, 1, dh=True, alpha=0.6, beta=0.2, gamma=50, nrad=2, mdis=12, sclip=nn3ss_h).std.Transpose().resize.Bicubic(src_top=-.5, src_left=-.5)
fsrcnnx = vdf.fsrcnnx_upscale(y, 3840, 2160, shader_file=r'C:\Users\User\AppData\Roaming\mpv.net\shaders\FSRCNNX_x2_56-16-4-1.glsl', downscaler=None, upscaled_smooth=nn3ssaa, profile='zastin')
#Clamp AA
aa_w = lvf.aa.clamp_aa(fsrcnnx, nn3ssaa, ee3ss_h, strength=1.3).resize.Bicubic(1920,1080).fmtc.bitdepth(bits=16)
aa_m = lvf.aa.clamp_aa(fsrcnnx, nn3ssaa, ee3ss_h, strength=2.8).resize.Bicubic(1920,1080).fmtc.bitdepth(bits=16)
#Mask AA
aa_ret    = kgf.retinex_edgemask(y.fmtc.bitdepth(bits=16), sigma=2)
aa_w_mask = core.std.Inflate(aa_ret.rgvs.RemoveGrain(3).rgvs.RemoveGrain(3).rgvs.RemoveGrain(11)).std.Deflate()
aa_m_mask = core.std.Binarize(aa_ret, 16000,0).rgvs.RemoveGrain(3).rgvs.RemoveGrain(3).rgvs.RemoveGrain(11)
#Merge AA
aa_w_merge = core.std.MaskedMerge(y.fmtc.bitdepth(bits=16), aa_w, aa_w_mask)
aa_m_merge = core.std.MaskedMerge(aa_w_merge, aa_m, aa_m_mask)
#Join AA
aa = join([aa_m_merge.fmtc.bitdepth(bits=32), u, v])
#Skip AA
aa = lvf.rfs(aa, src32, ranges=skip_aa)
###

###Denoising
smdg = hvf.SMDegrain(aa.fmtc.bitdepth(bits=16), tr=2, thSAD=200, blksize=8, overlap=4, contrasharp=False, RefineMotion=True, chroma=False, pel=1, plane=0).fmtc.bitdepth(bits=32)
bm3d = eoe.dn.BM3D(aa, sigma=[3.3], profile='np', radius=2, block_step=2, bm_range=12, ps_range=6, CUDA=True, ref=smdg, chroma=False, Fast=False)
knlm = core.knlm.KNLMeansCL(bm3d, d=3, a=2, s=2, h=0.5, channels = 'UV', device_type='gpu')
uvdn = core.std.Merge(knlm, ccd(knlm, threshold=2, matrix=1), weight=3/5).fmtc.bitdepth(bits=16)
###

###Dehalo
dehalo = hvf.FineDehalo(uvdn, rx=2, ry=2, thmi=85, thma=185, thlimi=8, thlima=80, brightstr=1.2, darkstr=0)
###

###Debanding
db_d3kdb  = dbs.dumb3kdb(dehalo, 17, [30,30,30])
db_d3kdb2 = dbs.dumb3kdb(dehalo, 5, [35,35,35])
db_d3kdb3 = dbs.dumb3kdb(dehalo, 24, [40,40,40])
db_mix    = core.average.Mean([db_d3kdb, db_d3kdb2, db_d3kdb3])
db_mask   = lvf.mask.detail_mask_neo(dehalo, sigma=1.8, detail_brz=0.009, lines_brz=0.018).rgvs.RemoveGrain(11).std.Inflate().std.Inflate()
db        = core.std.MaskedMerge(db_mix, dehalo, db_mask)
#Regrain
grain = adptvgrnMod(db, strength=[0.24,0.08], size=1.2, luma_scaling=20, sharp=60, static=True, fade_edges=True, show_mask=False, seed=8888)
###

###Outputs
#yt, ut, vt = split(src32.fmtc.bitdepth(bits=32)) #for testing
#yt2, ut2, vt2 = split(knlm.fmtc.bitdepth(bits=32)) #for testing
#yt3, ut3, vt3 = split(uvdn.fmtc.bitdepth(bits=32)) #for testing
#
out = core.fmtc.bitdepth(grain, bits=10)#.std.Invert()
#cut  = out[21412:22009] #for testing
#
out.set_output()
###