import vapoursynth as vs
import sys
import math
import havsfunc as haf
import mvsfunc as mvf

core = vs.get_core(threads = 12)
core.max_cache_size = 8000

def scale(val, bits):
		return val * ((1 << bits) - 1) // 255

def Levels(clip, input_low, gamma, input_high, output_low, output_high, planes=None, coring=True):
    core = vs.get_core()
    funcName = 'Levels'

    if not isinstance(clip, vs.VideoNode):
        raise TypeError(funcName + ': \"clip\" must be a clip!')

    if planes is None:
        planes = list(range(clip.format.num_planes))

    bits = clip.format.bits_per_sample
    neutral = 1 << (bits - 1)
    peak = (1 << bits) - 1

    gamma = 1 / gamma
    divisor = input_high - input_low + (input_high == input_low)

    tvLow = haf.scale(16, bits)
    tvHigh = [haf.scale(235, bits), haf.scale(240, bits)]
    scaleUp = peak / haf.scale(235 - 16, bits)
    scaleDown = haf.scale(235 - 16, bits) / peak

    isGray = clip.format.color_family == vs.GRAY
    chroma_planes = []
    if not isGray:
        if 1 in planes:
            chroma_planes.append(1)
        if 2 in planes:
            chroma_planes.append(2)

    def get_lut1(x):
        p = ((x - tvLow) * scaleUp - input_low) / divisor if coring else (x - input_low) / divisor
        p = min(max(p, 0), 1) ** gamma * (output_high - output_low) + output_low
        return min(max(math.floor(p * scaleDown + tvLow + 0.5), tvLow), tvHigh[0]) if coring else min(max(math.floor(p + 0.5), 0), peak)

    def get_lut2(x):
        q = math.floor((x - neutral) * (output_high - output_low) / divisor + neutral + 0.5)
        return min(max(q, tvLow), tvHigh[1]) if coring else min(max(q, 0), peak)

    if 0 in planes:
        clip = core.std.Lut(clip, planes=[0], function=get_lut1)
    if chroma_planes:
        clip = core.std.Lut(clip, planes=chroma_planes, function=get_lut2)

    return clip

def ringmask2(clip, rx=None, amp=None, thY1=0, thY2=255, mode=1):
    w = clip.width
    h = clip.height
    bits = clip.format.bits_per_sample
    peak = (1 << bits) - 1
    ry = rx

    thY1 = scale(thY1, bits)
    thY2 = scale(thY2, bits)

    smooth = core.fmtc.resample(clip, haf.m4(w / rx), haf.m4(h / ry), kernel='bicubic',fulls=False, fulld=False).fmtc.resample(w, h, kernel='bicubic', a1=1, a2=0,fulls=False, fulld=False)
    smoother = core.fmtc.resample(clip, haf.m4(w / rx), haf.m4(h / ry), kernel='bicubic',fulls=False, fulld=False).fmtc.resample(w, h, kernel='bicubic', a1=1.5, a2=-0.25,fulls=False, fulld=False)

    expr = 'x y - {amp} *'.format(amp=amp) if mode == 1 else 'y x - {amp} *'.format(amp=amp)
    mask = core.std.Expr([smooth, smoother], [expr]).fmtc.bitdepth(bits=bits, fulls=True, fulld=True, dmode=1)

    limitexpr = 'x {thY1} < 0 x {thY2} >= {peak} x ? ?'.format(thY1=thY1, thY2=thY2, peak=peak)
    mask = core.std.Expr([mask], [limitexpr])

    return mask

def dr_only(clp,strength1,strength2,dering_line_mask2):
	bits = clp.format.bits_per_sample
	multiple = scale(1, bits)
	x = clp.width
	y = clp.height
	x  = round((x / 4.0) * 4)
	y  = round((y / 4.0) * 4)
	x2 = x // 2
	y2 = y // 2
	xs = x * 3 // 2
	ys = y * 3 // 2
	upsampling_args = dict(kernel='lanczos', taps=4, noring=True)
	downsampling_args = dict(kernel='spline36')

	ps1 = core.std.MaskedMerge(clp, core.rgvs.RemoveGrain(clp, [20]), Levels(dering_line_mask2, scale(60, bits), 3.0, scale(140, bits), 0, scale(strength1, bits)))
	halo = core.fmtc.resample(ps1, x2, y2, kernel='bicubic',fulls=False, fulld=False).fmtc.resample(x, y, kernel='bicubic', a1=1, a2=0,fulls=False, fulld=False).fmtc.bitdepth(bits=bits,fulls=False, fulld=False)
	halo_mask1 = core.std.Expr([core.std.Maximum(ps1), core.std.Minimum(clp)], ['x y -'])
	halotomask = core.std.Expr([core.std.Maximum(halo), core.std.Minimum(halo)], ['x y -'])
	halotomask2 = core.std.Expr([halotomask, halo_mask1], ['y {multiple} / x {multiple} / - y {multiple} / 0.001 + / 255 * 50 - y {multiple} / 256 + 512 / 0.5 + * {multiple} *'.format(multiple=multiple)])
	pshalo = core.std.MaskedMerge(halo, ps1, halotomask2)
	pshalo2 = haf.Resize(ps1, xs, ys, bits=bits, **upsampling_args)
	pshalo2 = core.std.Expr([pshalo2, core.std.Maximum(pshalo).fmtc.resample(xs, ys, kernel='bicubic').fmtc.bitdepth(bits=bits)], ['x y min'])
	pshalo2 = core.std.Expr([pshalo2, core.std.Minimum(pshalo).fmtc.resample(xs, ys, kernel='bicubic').fmtc.bitdepth(bits=bits)], ['x y max'])
	pshalo2 = haf.Resize(pshalo2, x, y, bits=bits, **downsampling_args)
	ps2 = core.std.Expr([ps1, pshalo2], ['x y < x x y - {} * - x x y - 1.0 * - ?'.format(strength2)])

	return ps2

#OKE:INPUTFILE
a="00001.m2ts"
src8 = core.lsmas.LWLibavSource(a, threads=1)
src16 = core.fmtc.bitdepth(src8,bits=16)

Y = core.fmtc.resample(src16,960,540,sx=-0.5)
down = core.std.ShufflePlanes([Y,src16], [0,1,2], vs.YUV)
nr16y = core.knlm.KNLMeansCL(src16, d=1, a=3, s=2, wmode=3, h=2.0, device_type="GPU")
nr16uv = core.knlm.KNLMeansCL(down, d=1, a=2, s=2, channels="YUV", wmode=3, h=1.2, device_type="GPU")
nr16 = core.std.ShufflePlanes([nr16y,nr16uv], [0,1,2], vs.YUV)
noise16 = core.std.MakeDiff(src16, nr16)
pre16 = core.std.ShufflePlanes([nr16y,src16],[0,1,2],vs.YUV)
pre8 = mvf.Depth(pre16, depth=8, dither=5)

eemask = core.tcanny.TCanny(src8,sigma=0.8,mode=1,op=2,gmmax=255, planes=[0,1,2])
eemask_u = core.std.ShufflePlanes(eemask, [1,1,1], vs.YUV).fmtc.resample(1920,1080,sx=0.25,css="420")
eemask_v = core.std.ShufflePlanes(eemask, [2,2,2], vs.YUV).fmtc.resample(1920,1080,sx=0.25,css="420")
weighted = core.std.Expr([eemask,eemask_u,eemask_v],["x 128 * y + z +",""],vs.YUV420P16)

nrmasks = core.std.Binarize(weighted, 4800, 0)
dmaskb = core.std.Binarize(weighted, 2400, 0)
dmasks = core.std.Binarize(weighted, 3600, 0)
dmask = core.misc.Hysteresis(dmasks, dmaskb, 0)

aamask = core.std.Binarize(weighted, 6400, 0)
aamask = core.std.Maximum(aamask,0).std.Maximum(0).std.Minimum(0).std.Minimum(0)
aamask = core.rgvs.RemoveGrain(aamask,[20,0])

Y1080 = core.std.ShufflePlanes(pre16, 0, vs.GRAY)
U1080 = core.std.ShufflePlanes(pre8, 1, vs.GRAY).fmtc.resample(1920,1080,sx=0.25,kernel="lanczos")
V1080 = core.std.ShufflePlanes(pre8, 2, vs.GRAY).fmtc.resample(1920,1080,sx=0.25,kernel="lanczos")

textmask = core.std.Expr([Y1080,U1080,V1080],"x 57600 >= y 32768 - abs 770 <= and z 32768 - abs 770 <= and 65535 0 ?")
textmask_mod = core.std.Minimum(textmask).std.Minimum().std.Maximum().std.Maximum()
textmask = core.std.Expr([textmask, textmask_mod], "x y -")
textmask = core.std.ShufflePlanes([textmask,src16],[0,1,2],vs.YUV)
textmask = core.std.Maximum(textmask,0).std.Maximum(0).std.Maximum(0).rgvs.RemoveGrain([20,0]).rgvs.RemoveGrain([20,0])
aamask = core.std.Expr([aamask, textmask, weighted], ["z 10000 > x y - x ?",""], vs.YUV420P16)

luma = core.std.ShufflePlanes(pre8, 0, vs.YUV).resize.Bilinear(format=vs.YUV420P8)
eemask = core.std.Expr([eemask,luma],["x x * 60 * 255 y - dup * 0.4 * + 60000 min","x x * 80 * 255 y - dup * 0.6 * + 60000 min"],vs.YUV420P16)
eemask = core.std.Maximum(core.std.Maximum(eemask,planes=[0,1,2]),planes=[0,1,2])
eemask = core.rgvs.RemoveGrain(eemask, [20,11]).rgvs.RemoveGrain([20,11])

nrmaskg = core.tcanny.TCanny(pre8,sigma=1.0,t_h=8.0,t_l=1.2,op=2,planes=0)
nrmaskb = core.tcanny.TCanny(pre8,sigma=1.4,t_h=9.0,t_l=1.5,op=2,planes=0)

nrmask = core.std.Expr([nrmaskg,nrmaskb,nrmasks,pre8,dmask],["a 20 < 65535 a 64 < x 257 * b max a 128 < y 257 * z ? ? ?","",""],vs.YUV420P16)
nrmask = core.std.Maximum(nrmask,0).std.Maximum(0).std.Minimum(0)
nrmask = core.rgvs.RemoveGrain(core.rgvs.RemoveGrain(nrmask,[20,0]),[20,0])

w       = 1920
h       = 1080

aa_clip = core.std.ShufflePlanes(src16, 0, vs.GRAY)
aa_clip = core.eedi2.EEDI2(aa_clip, field=1, mthresh=10, lthresh=20, vthresh=20, maxd=24, nt=50)
aa_clip = core.fmtc.resample(aa_clip, w, h, sy=[-0.5,-1]).std.Transpose()
aa_clip = core.eedi2.EEDI2(aa_clip, field=1, mthresh=10, lthresh=20, vthresh=20, maxd=24, nt=50)
aa_clip = core.fmtc.resample(aa_clip, h, w, sy=[-0.5,-1]).std.Transpose()
edge = core.std.ShufflePlanes([aa_clip, src16], [0,1,2], vs.YUV)
aaed = core.std.MaskedMerge(src16, edge, aamask, [0,1,2], True)

preblur = core.std.MakeDiff(aaed,noise16)

dbed = core.f3kdb.Deband(preblur,12,48,64,64,0,0,output_depth=16)
dbed = core.f3kdb.Deband(dbed,24,32,48,48,0,0,output_depth=16)
dbed = mvf.LimitFilter(dbed, preblur, thr=0.6, thrc=0.8, elast=1.5)
dbed = core.std.MaskedMerge(dbed, preblur, nrmask, first_plane=1)

dif = core.std.MakeDiff(dbed, core.rgvs.RemoveGrain(dbed, 20))
sharp = core.std.MergeDiff(dbed, dif)
sharped = core.std.MaskedMerge(sharp, dbed, eemask, [0,1,2], False)

noise16 = core.std.Expr(noise16,["x 32768 - 1.05 * 32768 +",""])
nullclip = core.std.Expr(src16,["32768",""])
nrweight = core.std.Expr(pre8, ["x 64 - 0 max dup * 4 * ",""], vs.YUV420P16)
noise16 = core.std.MaskedMerge(noise16,nullclip,nrweight,0,True)
res = core.std.MergeDiff(sharped,noise16)

ringmask = ringmask2(src16, 1.5, 50)
ringmask = core.std.Expr([ringmask, luma, textmask],["y 16 > y 235 < and x z - 0 ?",""])
gamma=3.0
pre = core.std.Levels(res, 4096, 60160, 1/gamma, 4096, 60160)
dr = dr_only(pre,strength1=0,strength2=0,dering_line_mask2=ringmask)
dr = core.std.Levels(dr, 4096, 60160, gamma, 4096, 60160)
res = core.std.MaskedMerge(res, dr, ringmask, 0, True)
res = core.std.SetFrameProp(res, prop="_ChromaLocation", intval = 0)

#OKE:DEBUG
Debug = 1
if Debug:
	res=core.std.Interleave([src16, res])
	res=mvf.ToRGB(res,full=False,depth=8)
else: res = core.fmtc.bitdepth(res,bits=10)

res.set_output()