import vapoursynth
import asyncio
import gc
core = vapoursynth.core
core.max_cache_size=1024
loop = asyncio.get_event_loop()
ar = 1920 / 1080



def getw(h, only_even=True):
    w = h * ar
    w = int(round(w))
    if only_even:
        w = w // 2 * 2
    return w



async def run(src):
    # change format to GrayS with bitdepth 32 for descale
    src = src[0]
    matrix_s = '709' if src.format.color_family == vapoursynth.RGB else None
    src_luma32 = core.resize.Point(src, format=vapoursynth.YUV444PS, matrix_s=matrix_s)
    src_luma32 = core.std.ShufflePlanes(src_luma32, 0, vapoursynth.GRAY)
    src_luma32 = core.std.Cache(src_luma32)
    # descale each individual frame
    resizer = core.resize.Bilinear
    upscaler = core.resize.Bilinear
    clip_list = []
    for h in range(500, 2000 + 1):
        clip_list.append(resizer(src_luma32, getw(h), h))
    full_clip = core.std.Splice(clip_list, mismatch=True)
    full_clip = upscaler(full_clip, getw(src.height), src.height)
    if ar != src.width / src.height:
        src_luma32 = upscaler(src_luma32, getw(src.height), src.height)
    expr_full = core.std.Expr([src_luma32 * full_clip.num_frames, full_clip], 'x y - abs dup 0.015 > swap 0 ?')
    full_clip = core.std.CropRel(expr_full, 5, 5, 5, 5)
    full_clip = core.std.PlaneStats(full_clip)
    full_clip = core.std.Cache(full_clip)
    try:
        tasks_pending = set()
        futures = {}
        vals = []
        full_clip_len = len(full_clip)
        for frame_index in range(len(full_clip)):
            print(f"{frame_index+1}/{full_clip_len}", end="\r")
            fut = asyncio.ensure_future(asyncio.wrap_future(full_clip.get_frame_async(frame_index)))
            tasks_pending.add(fut)
            futures[fut] = frame_index
            while len(tasks_pending) >= core.num_threads + 2:
                tasks_done, tasks_pending = await asyncio.wait(tasks_pending, return_when=asyncio.FIRST_COMPLETED)
                vals += [(futures.pop(task), task.result().props.PlaneStatsAverage) for task in tasks_done]
        tasks_done, _ = await asyncio.wait(tasks_pending)
        vals += [(futures.pop(task), task.result().props.PlaneStatsAverage) for task in tasks_done]
        vals = [v for _, v in sorted(vals)]
    finally:
        gc.collect()



#windows
loop.run_until_complete(run(core.imwrif.Read(r"mpv-shot0001.png")))