import vapoursynth as vs
from vapoursynth import core
from vsutil import *
import awsmfunc as awf
import lvsfunc as lvf

b = 0
c = 0.5

height = 720 # yes that is a nice chad resolution
width = 1280

def nnedi3_rpow2(clip):
    clip = clip.std.Transpose().znedi3.nnedi3(0, 1, 0, 0, 4, 2).resize.Spline36(height=1920, src_top=.5)
    return clip.std.Transpose().znedi3.nnedi3(0, 1, 0, 0, 4, 2).resize.Spline36(height=1080, src_top=.5)

def descale(clip): return depth(depth(clip, 32).descale.Debicubic(width, height, b=b, c=c), 16)

def error_mask(clip, desc, thr1=2500, thr2=None, expand1=2, expand2=3, blur=3):
    start = 0

    error = desc.resize.Bicubic(1920, 1080, filter_param_a=b, filter_param_b=c)
    error = core.std.Expr([clip, error], 'x y - abs')
    error = iterate(error, core.std.Maximum, expand1)

    for x in range(expand1):
        error = error.std.Maximum(coordinates=[[1]*8, [0,1,0,1,1,0,1,0]][min(start%3,1)])
        start += 1

    if thr2 is not None:
        error = error.std.Binarize(thr2).misc.Hysteresis(error.std.Binarize(thr1))
    else:
        error = error.std.Binarize(thr1)

    for x in range(expand2):
        error = error.std.Maximum(coordinates=[[1]*8, [0,1,0,1,1,0,1,0]][min(start%3,1)])
        start += 1

    if blur:
        error = error.std.BoxBlur(hradius=blur, vradius=blur)

    error = error.std.Crop(5,5,5,5).std.AddBorders(5,5,5,5,65535)

    return error

def DescaleAAmod(clip):
    y,u,v = split(clip)
    desc = descale(y)
    resc = nnedi3_rpow2(desc)
    mask = y.std.Prewitt().std.Maximum().std.Binarize(24<<8).rgvs.RemoveGrain(20)
    error = error_mask(y, desc)
    y = y.std.MaskedMerge(resc, mask).std.MaskedMerge(y, error)
    return join([y,u,v])

src = lvf.misc.source("/mnt/data/samba/Uploads/Towa.no.Quon.6.Eternal.Quon.2011.1080p.BluRay.DTSHD-MA.h264.Remux-W33B.mkv") # load up the shitbox
src = depth(src, 16) # fucked if I know

descale_b = DescaleAAmod(src)

source = awf.FrameInfo(src, "Source") #screenshots for non-trumpability
descale_output = awf.FrameInfo(descale_b, "Descale") #same
comp = core.std.Interleave([source, descale_output]) #same
comp.set_output()