from vapoursynth import core
import vapoursynth as vs
import math

"""
A vapoursynth port of avisynth's bbmod, same usage, add 10 bit support.
"""

def bbmod2(c, cTop = None, cBottom = None, cLeft = None, cRight = None, thresh = 128, blur = 999):
    if blur < 1:
        raise TypeError("BalanceBordersMod2: Blur parameter must be > 0")
    if thresh < 1:
        raise TypeError("BalanceBordersMod2: Thresh parameter must be > 0")

    bits = c.format.bits_per_sample

    c = btb2(c, cTop, thresh, blur, bits).std.Transpose().std.FlipHorizontal() if cTop > 0 else core.std.Transpose(c).std.FlipHorizontal()
    c = btb2(c, cLeft, thresh, blur, bits).std.Transpose().std.FlipHorizontal() if cLeft > 0 else core.std.Transpose(c).std.FlipHorizontal()
    c = btb2(c, cBottom, thresh, blur, bits).std.Transpose().std.FlipHorizontal() if cBottom > 0 else core.std.Transpose(c).std.FlipHorizontal()
    c = btb2(c, cRight, thresh, blur, bits).std.Transpose().std.FlipHorizontal() if cRight > 0 else core.std.Transpose(c).std.FlipHorizontal()
    return c

def btb2(c, cTop, thresh, blur, bits):
    cWidth = c.width
    cHeight = c.height
    cTop = min(cTop, cHeight - 1)
    blurWidth = max(8,math.floor(cWidth/blur))

    if bits == 8:
        c2 = core.resize.Point(c, width = cWidth*2, height = cHeight*2)
        last = core.std.Crop(c2, 0, 0, cTop*2, cHeight*2-cTop*2-2)
        last = core.resize.Point(last, cWidth*2, cTop*2)
        referenceBlurChroma = core.std.Expr(last, "x 128 - abs 2 *").resize.Bicubic(blurWidth*2, cTop*2, filter_param_a = 1, filter_param_b = 0).resize.Bicubic(cWidth*2, cTop*2, filter_param_a = 1, filter_param_b = 0)
        referenceBlur = core.resize.Bicubic(last, blurWidth*2, cTop*2, filter_param_a = 1, filter_param_b = 0).resize.Bicubic(cWidth*2, cTop*2, filter_param_a = 1, filter_param_b = 0)

        original = core.std.Crop(c2, 0, 0, 0, c2.height-(cTop*2))
        last = core.resize.Bicubic(original, blurWidth*2, cTop*2, filter_param_a = 1, filter_param_b = 0)
        originalBlurChroma = core.std.Expr(last, "x 128 - abs 2 *").resize.Bicubic(blurWidth*2, cTop*2, filter_param_a = 1, filter_param_b = 0).resize.Bicubic(cWidth*2, cTop*2, filter_param_a = 1, filter_param_b = 0)
        originalBlur = core.resize.Bicubic(last, blurWidth*2, cTop*2, filter_param_a = 1, filter_param_b = 0).resize.Bicubic(cWidth*2, cTop*2, filter_param_a = 1, filter_param_b = 0)

        balancedChroma = core.std.Expr([original, originalBlurChroma, referenceBlurChroma],["","z y / 8 min 0.4 max x 128 - * 128 +", "z y / 8 min 0.4 max x 128 - * 128 +"])
        balancedLuma = core.std.Expr([balancedChroma, originalBlur, referenceBlur], ["z 16 - y 16 - / 8 min 0.4 max x 16 - * 16 +", "z y - x +", "z y - x +"])

        difference = core.std.MakeDiff(balancedLuma, original)

        tp = str(128+thresh)
        tm = str(128-thresh)
        difference = core.std.Expr(difference, "x "+tp+" > "+tp+" x ?")
        difference = core.std.Expr(difference, "x "+tm+" < "+tm+" x ?")

        last = core.std.MergeDiff(original, difference)
        return core.std.StackVertical([last,core.std.Crop(c2,0,0,cTop*2,c2.height-cTop*2-(cHeight-cTop)*2)]).resize.Point(cWidth,cHeight)

    if bits == 10:
        c2 = core.resize.Point(c, width = cWidth*2, height = cHeight*2)
        last = core.std.Crop(c2, 0, 0, cTop*2, cHeight*2-cTop*2-2)
        last = core.resize.Point(last, cWidth*2, cTop*2)
        referenceBlurChroma = core.std.Expr(last, "x 512 - abs 2 *").resize.Bicubic(blurWidth*2, cTop*2, filter_param_a = 1, filter_param_b = 0).resize.Bicubic(cWidth*2, cTop*2, filter_param_a = 1, filter_param_b = 0)
        referenceBlur = core.resize.Bicubic(last, blurWidth*2, cTop*2, filter_param_a = 1, filter_param_b = 0).resize.Bicubic(cWidth*2, cTop*2, filter_param_a = 1, filter_param_b = 0)

        original = core.std.Crop(c2, 0, 0, 0, c2.height-(cTop*2))
        last = core.resize.Bicubic(original, blurWidth*2, cTop*2, filter_param_a = 1, filter_param_b = 0)
        originalBlurChroma = core.std.Expr(last, "x 512 - abs 2 *").resize.Bicubic(blurWidth*2, cTop*2, filter_param_a = 1, filter_param_b = 0).resize.Bicubic(cWidth*2, cTop*2, filter_param_a = 1, filter_param_b = 0)
        originalBlur = core.resize.Bicubic(last, blurWidth*2, cTop*2, filter_param_a = 1, filter_param_b = 0).resize.Bicubic(cWidth*2, cTop*2, filter_param_a = 1, filter_param_b = 0)

        balancedChroma = core.std.Expr([original, originalBlurChroma, referenceBlurChroma],["","z y / 8 min 0.4 max x 512 - * 512 +", "z y / 8 min 0.4 max x 512 - * 512 +"])
        balancedLuma = core.std.Expr([balancedChroma, originalBlur, referenceBlur], ["z 64 - y 64 - / 8 min 0.4 max x 64 - * 64 +", "z y - x +", "z y - x +"])

        difference = core.std.MakeDiff(balancedLuma, original)

        tp = str(512+thresh)
        tm = str(512-thresh)
        difference = core.std.Expr(difference, "x "+tp+" > "+tp+" x ?")
        difference = core.std.Expr(difference, "x "+tm+" < "+tm+" x ?")

        last = core.std.MergeDiff(original, difference)
        return core.std.StackVertical([last,core.std.Crop(c2,0,0,cTop*2,c2.height-cTop*2-(cHeight-cTop)*2)]).resize.Point(cWidth,cHeight)