import vapoursynth as vs
import functools

def logic_andn_clip(clipa, clipb):
    return vs.get_core().std.Lut2(clipa=clipa, clipb=clipb,
                                   function=lambda x, y: x & (~y))

def logic_and_clip(clipa, clipb):
    return vs.get_core().std.Lut2(clipa=clipa, clipb=clipb,
                                   function=lambda x, y: x & y)

def make_nmasks(c1, c2, thr=10):
    core = vs.get_core()
    cm = functools.partial(core.comb.CombMask, cthresh=thr, mthresh=thr,
                           planes=[0, 1, 2])

    cmask = cm(c2)
    nmask1 = core.generic.Invert(clip=cmask, planes=[0, 1, 2]) # 255(combed) -> 0(combed)
    nmask2 = logic_andn_clip(cmask, cm(c1))
    return nmask1, nmask2

def flexible_merge(clip1, clip2, nmask1=None, nmask2=None, weight=0.5,
                     weight2=0.9, thr=10, exc=0, nmax=0.12, cmthr=10, diffthr=0,
                     diffthr2=0.90, show=0, plane=0) :
    core = vs.get_core()

    planes = range(clip1.format.num_planes)
    if len(planes) > 3:
        raise ValueError('flexible_merge: invalid the number of planes!')

    def select_noiseless_clip(n, f, clip1, clip2):
        if f[0].props.PlaneAverage < f[1].props.PlaneAverage:
            return clip1
        return clip2

    def select_clip_by_threshold(n, f, clip1, clip2, core, threshold,
                                    show=0, plane=0):
        diff = abs(f[0].props.PlaneAverage - f[1].props.PlaneAverage)
        clip = clip1 if diff < threshold else clip2
        if show:
            yuv = ["\nY: ", "\n\nU: ", "\n\n\nV: "]
            text = "{yuv} {n} selective {diff} >= {th}".format(
                yuv=yuv[plane], n=n, diff=diff, th=threshold)
            clip = core.text.Text(clip, text=text)
        return clip

    def noise_exclude(clip, clip1, clip2, mode=1, thr=0.12, cmthresh=10,
                        showi=0, plane=0):
        partial = functools.partial
        cm = partial(core.comb.CombMask, cthresh=cmthresh, planes=planes)
        show = showi == 3

        if show:
            clip1 = core.text.Text(clip1, text="1")
            clip2 = core.text.Text(clip2, text="2")

        noise_avg = [core.std.PlaneAverage(cm(clip1), plane=0),
                     core.std.PlaneAverage(cm(clip2), plane=0)]

        scbt = partial(select_clip_by_threshold, clip1=clip, core=core,
                       threshold=thr, show=show, plane=plane)

        if mode == 1:
            return core.std.FrameEval(clip, eval=partial(scbt, clip2=clip1),
                                       prop_src=noise_avg)

        snc = partial(select_noiseless_clip, clip1=clip1, clip2=clip2)
        clip_select = core.std.FrameEval(clip, eval=snc, prop_src=noise_avg)

        return core.std.FrameEval(clip, eval=partial(scbt, clip2=clip_select),
                                   prop_src=noise_avg)

    def exclude_overdiff(clip1, clip2) :
        def difference_pixel_with_thr(x, y) :
            return (abs(x - y) < diffthr) * 255

        clip = core.std.Lut2(clipa=clip1, clipb=clip2,
                             function=difference_pixel_with_thr)

        # greater value means both clips are closed to.
        avgc = core.std.PlaneAverage(clip, plane=0)

        zero = core.std.BlankClip(clip=clip, color=[0]*clip.format.num_planes)
        zero = core.std.PlaneAverage(zero, plane=0)

        scbt = functools.partial(select_clip_by_threshold, clip1=zero,
                                 clip2=clip, core=core, threshold=diffthr2,
                                 show=0, plane=plane)

        return core.std.FrameEval(clip, eval=scbt, prop_src=[avgc, zero])

    if !isinstance(nmask1, vs.VideoNode) or !isinstance(nmask2, vs.VideoNode):
        nmask1, nmask2 = make_nmasks(clip1, clip2, thr=thr)

    if diffthr > 0 :
        dmask = exclude_overdiff(clip1, clip2)
        nmask1 = logic_and_clip(dmask, nmask1)
        nmask2 = logic_and_clip(dmask, nmask2)
        if exc == 2 :
            exc = 1

    clip3 = core.std.Merge(clip1, clip2, weight)
    clip4 = core.std.Merge(clip1, clip2, weight2)
    clip = core.std.MaskedMerge(clip1, clip3, nmask1, planes=planes)
    clip = core.std.MaskedMerge(clip,  clip4, nmask2, planes=planes)

    if exc > 0 :# and diffthr == 0 :
        clip = noise_exclude(clip, clip1, clip2, mode = exc, thr = nmax,
                             cmthresh = cmthr, showi = show, plane=plane)

    return clip

def flexible_merge3(c1, c2, weight=0.5, weight2=0.5, thr=10, exc=0,
                      nmax=0.12, cmthr=10, diffthr_y=0, diffthr_uv=0,
                      diffthr2=0.95, show_y=0, show_u=0, show_v=0):
    core = vs.get_core()

    def extract_clip(clip):
        sp = core.std.ShufflePlanes
        return (sp(clips, planes=p, colorfamily=vs.GRAY) for p in range(3))

    y1, u1, v1 = extract_clip(c1)
    y2, u2, v2 = extract_clip(c2)

    nmask1, nmask2 = make_nmasks(c1, c2, thr=thr)

    nmask1_y, nmask1_u, nmask1_v = extract_clip(nmask1)
    nmask2_y, nmask2_u, nmask2_v = extract_clip(nmask2)

    fm = functools.partial(flexible_merge, weight=weight, weight2=weight2,
                           thr=thr, exc=exc, diffthr2=diffthr2)
    y3 = fm(y1, y2, nmask1=nmask1_y, nmask2=nmask2_y, nmax=nmax, cmthr=cmthr,
            diffthr=diffthr_y, show = show_y, plane=0 )
    u3 = fm(u1, u2, nmask1=nmask1_u, nmask2=nmask2_u, nmax=nmax*2,
            cmthr=cmthr//2, diffthr=diffthr_uv, show = show_u, plane=1 )
    v3 = fm(v1, v2, nmask1=nmask1_v, nmask2=nmask2_v, nmax=nmax*2,
            cmthr=cmthr//2, diffthr=diffthr_uv, show = show_v, plane=2 )

    return core.std.ShufflePlanes(clips=[y3, u3, v3], planes=[0, 0, 0],
                                   colorfamily=vs.YUV)