import vapoursynth as vs

class InvalidArgument(Exception):
    def __init__(self, value):
        self.value = value
    def __str__(self):
        return repr(self.value)

class MonoS(object):
    Gcutoff = 70
    Ggain = 0.75

    #x < cutoff ? 0 : x * gain * (256+x) / 256
    def f16(x):
        if (x < Gcutoff):
            return 0
        else:
            result = int(float(x) * Ggain * ((0xFFFF+1)+x) / (0xFFFF+1))
            if result > 0xFFFF:
                return 0xFFFF
            else:
                return result
    def f8(x):
        if (x < Gcutoff):
            return 0
        else:
            result = int(float(x) * Ggain * ((0xFF+1)+x) / (0xFF+1))
            if result > 0xFF:
                return 0xFF
            else:
                return result

    def Luma16(x):
        p = x << 4
        if (p & (0xffff + 1)):
            return (0xffff - (p&0xffff))
        else:
            return p&0xffff

    def MaskDetail16(clip, final_width, final_height, RGmode=3, cutoff=17990, gain=0.75, expandN=2, inflateN=1, blur_more=False, src_left=0, src_top=0, src_width=0, src_height=0, kernel="bilinear", invkstaps = 4, taps=4):

        core = vs.get_core()

        global Gcutoff
        global Ggain

        Ggain = gain
        Gcutoff = cutoff
        clip = core.std.ShufflePlanes(clip, planes=0, colorfamily=vs.GRAY)
        temp1 = core.fmtc.resample(clip, final_width, final_height, kernel=kernel, invks=True, invkstaps=invkstaps, taps=taps).fmtc.resample(clip.width, clip.height, kernel=kernel, taps=taps)
        diff = core.std.MakeDiff(clip, temp1, 0)

        initial_mask = core.std.Lut(diff, function=MonoS.Luma16).rgvs.RemoveGrain(mode=[RGmode]).std.Lut(function=MonoS.f16)

        expanded = initial_mask
        for i in range(0, expandN):
            expanded = core.generic.Maximum(expanded, planes=[0])

        inflated = expanded
        for i in range(0, inflateN):
            inflated = core.generic.Inflate(inflated, planes=[0])

        final = core.fmtc.resample(inflated, final_width, final_height, src_left, src_top, src_width, src_height, taps=taps)

        if(blur_more):
            final = core.rgvs.RemoveGrain(final, mode=[12,0,0])

        final = core.std.ShufflePlanes(final, planes=0, colorfamily=vs.GRAY)
        return final

    def MaskDetail8(clip, final_width, final_height, RGmode=3, cutoff=70, gain=0.75, expandN=2, inflateN=1, blur_more=False, src_left=0, src_top=0, src_width=0, src_height=0, kernel="bilinear", invkstaps = 4, taps=4):

        core = vs.get_core()

        global Gcutoff
        global Ggain

        Ggain = gain
        Gcutoff = cutoff
        clip = core.std.ShufflePlanes(clip, planes=0, colorfamily=vs.GRAY)
        temp1 = core.fmtc.resample(clip, final_width, final_height, kernel=kernel, invks=True, invkstaps=invkstaps, taps=taps).fmtc.resample(clip.width, clip.height, kernel=kernel, taps=taps).fmtc.bitdepth(bits=8)
        diff = core.std.MakeDiff(clip.fmtc.bitdepth(bits=8), temp1, 0)

        initial_mask = core.hist.Luma(diff).rgvs.RemoveGrain(mode=[RGmode]).std.Lut(function=MonoS.f8)

        expanded = initial_mask
        for i in range(0, expandN):
            expanded = core.generic.Maximum(expanded, planes=[0])

        inflated = expanded
        for i in range(0, inflateN):
            inflated = core.generic.Inflate(inflated, planes=[0])

        final = core.fmtc.resample(inflated, final_width, final_height, src_left, src_top, src_width, src_height, taps=taps).fmtc.bitdepth(bits=8)

        if(blur_more):
            final = core.rgvs.RemoveGrain(final, mode=[12,0,0])

        final = core.std.ShufflePlanes(final, planes=0, colorfamily=vs.GRAY)
        return final

    def MaskDetail(clip, final_width, final_height, RGmode=3, cutoff=0, gain=0.75, expandN=2, inflateN=1, blur_more=False, src_left=0, src_top=0, src_width=0, src_height=0, kernel="bilinear", invkstaps = 4, taps=4):
        if ((clip.format.bits_per_sample == 8) or (clip.format.bits_per_sample == 16)):
            if (clip.format.bits_per_sample == 8):
                if cutoff == 0:
                    cutoff = 70
                return MonoS.MaskDetail8(clip, final_width, final_height, RGmode=3, cutoff=cutoff, gain=0.75, expandN=2, inflateN=1, blur_more=False, src_left=0, src_top=0, src_width=0, src_height=0, kernel="bilinear", invkstaps = 4, taps=4)
            else:
                if cutoff == 0:
                    cutoff = 17990
                return MonoS.MaskDetail16(clip, final_width, final_height, RGmode=3, cutoff=cutoff, gain=0.75, expandN=2, inflateN=1, blur_more=False, src_left=0, src_top=0, src_width=0, src_height=0, kernel="bilinear", invkstaps = 4, taps=4)
        else:
            raise InvalidArgument('Input clip must be 8 or 16 bit')