'''
    2D version of 1st-level algorithm is a combination of frame_blobs, intra_blob, and comp_P: optional raster-to-vector conversion.
    intra_blob recursively evaluates each blob for two forks of extended internal cross-comparison and sub-clustering:
    
    der+: incremental derivation cross-comp in high-variation edge areas of +vg: positive deviation of gradient triggers comp_g, 
    rng+: incremental range cross-comp in low-variation flat areas of +v--vg: positive deviation of negated -vg triggers comp_r.
    Each adds a layer of sub_blobs per blob.  
    Please see diagram: https://github.com/boris-kz/CogAlg/blob/master/frame_2D_alg/Illustrations/intra_blob_2_fork_scheme.png
    
    Blob structure, for all layers of blob hierarchy:
    root_dert__,  
    Dert = I, iDy, iDx, G, Dy, Dx, M, S (area), Ly (vertical dimension)
    # I: input, (iDy, iDx): angle of input gradient, G: gradient, (Dy, Dx): vertical and lateral Ds, M: match  
    sign, 
    box,  # y0, yn, x0, xn
    dert__,  # box of derts, each = i, idy, idx, g, dy, dx, m
    stack_[ stack_params, Py_ [(P_params, dert_)]]: refs down blob formation tree, in vertical (horizontal) order
    # next fork:
    fcr,  # flag comp rng, also clustering criterion in dert and Dert: g in der+ fork, i+m in rng+ fork? 
    fig,  # flag input is gradient
    rdn,  # redundancy to higher layers
    rng,  # comp range
    sub_layers  # [sub_blobs ]: list of layers across sub_blob derivation tree
                # deeper layers are nested, multiple forks: no single set of fork params?
'''
from collections import deque, defaultdict
from class_cluster import ClusterStructure, NoneType
from class_bind import AdjBinder
from frame_blobs_yx import assign_adjacents
from intra_comp_g import comp_g, comp_r
from itertools import zip_longest
from class_stream import BlobStreamer
from utils import pairwise
import numpy as np
# from comp_P_draft import comp_P_blob

# filters, All *= rdn:
ave = 50  # fixed cost per dert, from average m, reflects blob definition cost, may be different for comp_a?
aveB = 50  # fixed cost per intra_blob comp and clustering

class CDeepP(ClusterStructure):
    I = int
    G = int
    Dy = int
    Dx = int
    M = int
    iDy = int
    iDx = int
    L = int
    x0 = int
    sign = NoneType

class CDeepStack(ClusterStructure):
    I = int
    G = int
    Dy = int
    Dx = int
    M = int
    iDy = int
    iDx = int
    S = int
    Ly = int
    y0 = int
    Py_ = list
    blob = object
    down_connect_cnt = int
    sign = NoneType

class CDeepBlob(ClusterStructure):
    Dert = dict
    box = list
    stack_ = list
    sign = NoneType
    open_stacks = int
    root_dert__ = object
    dert__ = object
    mask = object
    adj_blobs = list
    fopen = bool
    margin = list
    fcr = bool
    fig = bool
    rdn = float
    rng = int
    Ls = int  # for visibility and next-fork rdn
    sub_layers = list

# --------------------------------------------------------------------------------------------------------------
# functions, ALL WORK-IN-PROGRESS:

def intra_blob(blob, rdn, rng, fig, fcr, **kwargs):  # recursive input rng+ | der+ cross-comp within blob
    # fig: flag input is g | p, fcr: flag comp over rng+ | der+
    if kwargs.get('render', None) is not None:  # stop rendering sub-blobs when blob is too small
        if blob.Dert['S'] < 100:
            kwargs['render'] = False

    spliced_layers = []  # to extend root_blob sub_layers
    ext_dert__, ext_mask = extend_dert(blob)
    if fcr:
        dert__, mask = comp_r(ext_dert__, fig, fcr, ext_mask)  # -> m sub_blobs
    else:
        dert__, mask = comp_g(ext_dert__, ext_mask)  # -> g sub_blobs:

    if dert__[0].shape[0] > 2 and dert__[0].shape[1] > 2 and False in mask:  # min size in y and x, least one dert in dert__
        sub_blobs = cluster_derts(dert__, mask, ave * rdn, fcr, fig, **kwargs)
        # fork params:
        blob.fcr = fcr
        blob.fig = fig
        blob.rdn = rdn
        blob.rng = rng
        blob.Ls = len(sub_blobs)  # for visibility and next-fork rdn
        blob.sub_layers = [sub_blobs]  # 1st layer of sub_blobs

        for sub_blob in sub_blobs:  # evaluate for intra_blob comp_g | comp_r:

            G = blob.Dert['G']; adj_G = blob.adj_blobs[2]
            borrow = min(abs(G), abs(adj_G) / 2)  # or adjacent M if negative sign?

            if sub_blob.sign:
                if sub_blob.Dert['M'] - borrow > aveB * rdn:  # M - (intra_comp value lend to edge blob)
                    # comp_r fork:
                    blob.sub_layers += intra_blob(sub_blob, rdn + 1 + 1 / blob.Ls, rng * 2, fig=fig, fcr=1, **kwargs)
                # else: comp_P_
            elif sub_blob.Dert['G'] + borrow > aveB * rdn:  # G + (intra_comp value borrow from flat blob)
                # comp_g fork:
                blob.sub_layers += intra_blob(sub_blob, rdn + 1 + 1 / blob.Ls, rng=rng, fig=1, fcr=0, **kwargs)
            # else: comp_P_

        spliced_layers = [spliced_layers + sub_layers for spliced_layers, sub_layers in
                          zip_longest(spliced_layers, blob.sub_layers, fillvalue=[])]
    return spliced_layers


def cluster_derts(dert__, mask, Ave, fcr, fig, render=False):  # similar to frame_to_blobs

    if fcr:  # comp_r output;  form clustering criterion:
        if fig:
            crit__ = dert__[0] + dert__[6] - Ave  # eval by i + m, accum in rng; dert__[:,:,0] if not transposed
        else:
            crit__ = Ave - dert__[3]  # eval by -g, accum in rng
    else:  # comp_g output
        crit__ = dert__[6] - Ave  # comp_g output eval by m, or clustering is always by m?

    root_dert__ = dert__  # derts after the comps operation, which is the root_dert__
    dert__ = [*zip(*dert__)]  # transpose dert__ into shape [y, params, x]

    sub_blobs = []  # from form_blob:
    stack_ = deque()  # buffer of running vertical stacks of Ps
    stack_binder = AdjBinder(CDeepStack)
    if render:
        streamer = BlobStreamer(CDeepBlob, crit__, mask)

    if render:
        streamer = BlobStreamer(CDeepBlob, crit__, mask)
    for y, dert_ in enumerate(dert__):  # in height, first and last row are discarded;  print(f'Processing intra line {y}...')
        # if False in mask[i]:  # [y,x,params], there is at least one dert in line
        P_binder = AdjBinder(CDeepP)  # binder needs data about clusters of the same level
        P_ = form_P_(zip(*dert_), crit__[y], mask[y], P_binder)  # horizontal clustering, adds a row of Ps
        if render:
            render = streamer.update_blob_conversion(y, P_)  # if return False, stop rendering
        P_ = scan_P_(P_, stack_, root_dert__, sub_blobs, P_binder)  # vertical clustering, adds up_connects per P and down_connect_cnt per stack
        stack_ = form_stack_(P_, root_dert__, sub_blobs, y)
        stack_binder.bind_from_lower(P_binder)

    while stack_:  # frame ends, last-line stacks are merged into their blobs:
        form_blob(stack_.popleft(), root_dert__, sub_blobs)

    blob_binder = AdjBinder(CDeepBlob)
    blob_binder.bind_from_lower(stack_binder)
    assign_adjacents(blob_binder)  # add adj_blobs to each blob
    # sub_blobs = find_adjacent(sub_blobs)

    if render:  # rendering mode after blob conversion
        streamer.end_blob_conversion(y)

    return sub_blobs


# clustering functions:
# -------------------------------------------------------------------------------------------------------------------

def form_P_(dert_, crit_, mask_, binder):  # segment dert__ into P__, in horizontal ) vertical order

    P_ = deque()  # row of Ps
    sign_ = crit_ > 0
    x0 = 0
    try:
        while mask_[x0]:  # skip until not masked
            next(dert_)
            x0 += 1
    except IndexError:
        return P_  # the whole line is masked, return an empty P

    I, iDy, iDx, G, Dy, Dx, M, L = *next(dert_), 1  # initialize P params
    _sign = sign_[x0]
    _mask = mask_[x0]  # mask bit per dert

    for x, (i, idy, idx, g, dy, dx, m) in enumerate(dert_, start=x0+1):  # loop left to right in each row of derts
        mask = mask_[x]
        if ~mask:  # current dert is not masked
            sign = sign_[x]
            if ~_mask and sign != _sign:  # prior dert is not masked and sign changed
                # pack P
                P = CDeepP(I=I, G=G, Dy=Dy, Dx=Dx, M=M, iDy=iDy, iDx=iDx, L=L,x0=x0, sign=_sign)
                P_.append(P)
                # initialize P params:
                I, iDy, iDx, G, Dy, Dx, M, L, x0 = 0, 0, 0, 0, 0, 0, 0, 0, x
            elif _mask:
                I, iDy, iDx, G, Dy, Dx, M, L, x0 = 0, 0, 0, 0, 0, 0, 0, 0, x
        # current dert is masked
        elif ~_mask:  # prior dert is not masked
            # pack P
            P = CDeepP(I=I, G=G, Dy=Dy, Dx=Dx, M=M, iDy=iDy, iDx=iDx, L=L, x0=x0, sign=_sign)
            P_.append(P)
            # initialize P params: (redundant)
            # I, iDy, iDx, G, Dy, Dx, M, L, x0 = 0, 0, 0, 0, 0, 0, 0, 0, x + 1

        if ~mask:  # accumulate P params:
            I += i
            iDy += idy
            iDx += idx
            G += g
            Dy += dy
            Dx += dx
            M += m
            L += 1
            _sign = sign  # prior sign
        _mask = mask

    if ~_mask:  # terminate and pack last P in a row if prior dert is unmasked
        P = CDeepP(I=I, G=G, Dy=Dy, Dx=Dx, M=M, iDy=iDy, iDx=iDx, L=L, x0=x0, sign=_sign)
        P_.append(P)

    for _P, P in pairwise(P_):
        if _P.x0 + _P.L == P.x0:  # check if Ps are adjacents
            binder.bind(_P, P)

    return P_


def scan_P_(P_, stack_, root_dert__, sub_blobs, binder):  # merge P into higher-row stack of Ps with same sign and x_coord overlap

    next_P_ = deque()  # to recycle P + up_connect_ that finished scanning _P, will be converted into next_stack_

    if P_ and stack_:  # if both input row and higher row have any Ps / _Ps left

        P = P_.popleft()  # load left-most (lowest-x) input-row P
        stack = stack_.popleft()  # higher-row stacks
        _P = stack.Py_[-1]  # last element of each stack is higher-row P
        up_connect_ = []  # list of same-sign x-overlapping _Ps per P

        while True:  # while both P_ and stack_ are not empty

            x0 = P.x0  # first x in P
            xn = x0 + P.L  # first x beyond P
            _x0 = _P.x0  # first x in _P
            _xn = _x0 + _P.L  # first x beyond _P

            if stack.G > 0:  # check for overlaps in 8 directions, else a blob may leak through its external blob
                if _x0 - 1 < xn and x0 < _xn + 1:  # x overlap between loaded P and _P
                    if P.sign == stack.sign:  # sign match
                        stack.down_connect_cnt += 1
                        up_connect_.append(stack)  # buffer P-connected higher-row stacks into P' up_connect_
                    else:
                        binder.bind(_P, P)

            else:  # -G, check for orthogonal overlaps only: 4 directions, edge blobs are more selective
                if _x0 < xn and x0 < _xn:  # x overlap between loaded P and _P
                    if P.sign == stack.sign:  # sign match
                        stack.down_connect_cnt += 1
                        up_connect_.append(stack)  # buffer P-connected higher-row stacks into P' up_connect_
                    else:
                        binder.bind(_P, P)

            if (xn < _xn or  # _P overlaps next P in P_
                    xn == _xn and stack.sign):  # sign taken accounted
                next_P_.append((P, up_connect_))  # recycle _P for the next run of scan_P_
                up_connect_ = []
                if P_:
                    P = P_.popleft()  # load next P
                else:  # terminate loop
                    if stack.down_connect_cnt != 1:  # terminate stack, merge it into up_connects' blobs
                        form_blob(stack, root_dert__, sub_blobs)
                    break
            else:  # no next-P overlap
                if stack.down_connect_cnt != 1:  # terminate stack, merge it into up_connects' blobs
                    form_blob(stack, root_dert__, sub_blobs)
                if stack_:  # load stack with next _P
                    stack = stack_.popleft()
                    _P = stack.Py_[-1]
                else:  # no stack left: terminate loop
                    next_P_.append((P, up_connect_))
                    break

    while P_:  # terminate Ps and stacks that continue at row's end
        next_P_.append((P_.popleft(), []))  # no up_connect
    while stack_:
        form_blob(stack_.popleft(), root_dert__, sub_blobs)  # down_connect_cnt always == 0

    return next_P_  # each element is P + up_connect_ refs


def form_stack_(P_, root_dert__, sub_blobs, y):  # Convert or merge every P into its stack of Ps, merge blobs

    next_stack_ = deque()  # converted to stack_ in the next run of scan_P_

    while P_:
        P, up_connect_ = P_.popleft()
        I, G, Dy, Dx, M, iDy, iDx, L, x0, s = P.unpack()
        xn = x0 + L  # next-P x0
        if not up_connect_:
            # initialize new stack for each input-row P that has no connections in higher row:
            blob = CDeepBlob(Dert=dict(I=0, G=0, Dy=0, Dx=0, M=0, iDy=0, iDx=0, S=0, Ly=0),
                             box=[y, x0, xn], stack_=[], sign=s, open_stacks=1)
            new_stack = CDeepStack(I=I, G=G, Dy=0, Dx=Dx, M=M, iDy=iDy, iDx=iDx, S=L, Ly=1,
                                   y0=y, Py_=[P], blob=blob, down_connect_cnt=0, sign=s)
            new_stack.hid = blob.id
            blob.stack_.append(new_stack)
        else:
            if len(up_connect_) == 1 and up_connect_[0].down_connect_cnt == 1:
                # P has one up_connect and that up_connect has one down_connect=P: merge P into up_connect stack:
                new_stack = up_connect_[0]
                new_stack.accumulate(I=I, G=G, Dy=Dy, Dx=Dx, M=M, iDy=iDy, iDx=iDx, S=L, Ly=1)
                new_stack.Py_.append(P)  # Py_: vertical buffer of Ps
                new_stack.down_connect_cnt = 0  # reset down_connect_cnt
                blob = new_stack.blob

            else:  # if > 1 up_connects, or 1 up_connect that has > 1 down_connect_cnt:
                blob = up_connect_[0].blob
                # initialize new_stack with up_connect blob:
                new_stack = CDeepStack(I=I, G=G, Dy=0, Dx=Dx, M=M, iDy=iDy, iDx=iDx, S=L, Ly=1,
                                       y0=y, Py_=[P], blob=blob, down_connect_cnt=0, sign=s)
                new_stack.hid = blob.id
                blob.stack_.append(new_stack)

                if len(up_connect_) > 1:  # merge blobs of all up_connects
                    if up_connect_[0].down_connect_cnt == 1:  # up_connect is not terminated
                        form_blob(up_connect_[0], root_dert__, sub_blobs)  # merge stack of 1st up_connect into its blob

                    for up_connect in up_connect_[1:len(up_connect_)]:  # merge blobs of other up_connects into blob of 1st up_connect
                        if up_connect.down_connect_cnt == 1:
                            form_blob(up_connect, root_dert__, sub_blobs)

                        if not up_connect.blob is blob:
                            merged_blob = up_connect.blob
                            I, G, Dy, Dx, M, iDy, iDx, S, Ly = merged_blob.Dert.values()
                            accum_Dert(blob.Dert, I=I, G=G, Dy=Dy, Dx=Dx, M=M, iDy=iDy, iDx=iDx, S=S, Ly=Ly)
                            blob.open_stacks += merged_blob.open_stacks
                            blob.box[0] = min(blob.box[0], merged_blob.box[0])  # extend box y0
                            blob.box[1] = min(blob.box[1], merged_blob.box[1])  # extend box x0
                            blob.box[2] = max(blob.box[2], merged_blob.box[2])  # extend box xn
                            for stack in merged_blob.stack_:
                                if not stack is up_connect:
                                    stack.blob = blob  # blobs in other up_connects are references to blob in the first up_connect.
                                    stack.hid = blob.id
                                    blob.stack_.append(stack)  # buffer of merged root stacks.
                            up_connect.blob = blob
                            up_connect.hid = blob.id
                            blob.stack_.append(up_connect)
                        blob.open_stacks -= 1  # overlap with merged blob.

        blob.box[1] = min(blob.box[1], x0)  # extend box x0
        blob.box[2] = max(blob.box[2], xn)  # extend box xn
        P.hid = new_stack.id  # assign higher cluster id for P
        next_stack_.append(new_stack)

    return next_stack_


def form_blob(stack, root_dert__, sub_blobs):  # increment blob with terminated stack, check for blob termination

    I, G, Dy, Dx, M, iDy, iDx, S, Ly, y0, Py_, blob, down_connect_cnt, sign = stack.unpack()
    accum_Dert(blob.Dert, I=I, G=G, Dy=Dy, Dx=Dx, M=M, iDy=iDy, iDx=iDx, S=S, Ly=Ly)
    # terminated stack is merged into continued or initialized blob (all connected stacks):

    blob.open_stacks += down_connect_cnt - 1  # incomplete stack cnt + terminated stack down_connect_cnt - 1: stack itself
    # open stacks contain Ps of a current row and may be extended with new x-overlapping Ps in next run of scan_P_
    if blob.open_stacks == 0:  # if number of incomplete stacks == 0
        # blob is terminated and packed in blob root:
        last_stack = stack
        y0, x0, xn = blob.box
        yn = last_stack.y0 + last_stack.Ly

        mask = np.ones((yn - y0, xn - x0), dtype=bool)  # mask box, then unmask Ps:
        for stack in blob.stack_:
            for y, P in enumerate(stack.Py_, start=stack.y0 - y0):
                x_start = P.x0 - x0
                x_stop = x_start + P.L
                mask[y, x_start:x_stop] = False

        fopen = 0  # flag: blob on frame boundary
        if x0 == 0 or xn == root_dert__[0].shape[1] or y0 == 0 or yn == root_dert__[0].shape[0]:
            fopen = 1

        blob.root_dert__ = root_dert__
        blob.box = (y0, yn, x0, xn)
        blob.dert__ = [derts[y0:yn, x0:xn] for derts in root_dert__]
        blob.mask = mask
        blob.adj_blobs = [[], 0, 0]
        blob.fopen = fopen

        sub_blobs.append(blob)


def extend_dert(blob):  # extend dert borders (+1 dert to boundaries)

    y0, yn, x0, xn = blob.box  # extend dert box:
    rY, rX = blob.root_dert__[0].shape  # higher dert size

    # determine pad size
    y0e = max(0, y0 - 1)
    yne = min(rY, yn + 1)
    x0e = max(0, x0 - 1)
    xne = min(rX, xn + 1)  # e is for extended

    # take ext_dert__ from part of root_dert__
    ext_dert__ = [derts[y0e:yne, x0e:xne] if derts is not None else None
                  for derts in blob.root_dert__]

    # pad mask: top, btm, left, right. 1 or 0 at boundaries
    mask = np.pad(blob.mask, ((y0 - y0e, yne - yn), (x0 - x0e, xne - xn)),
                  mode='constant', constant_values=True)

    return ext_dert__, mask


def accum_Dert(Dert: dict, **params) -> None:
    Dert.update({param: Dert[param] + value for param, value in params.items()})
[((148, 13, 148, 20), 'collections.deque', 'deque', ({}, {}), '()', False, 'from collections import deque, defaultdict\n'), ((149, 19, 149, 40), 'class_bind.AdjBinder', 'AdjBinder', ({(149, 29, 149, 39): 'CDeepStack'}, {}), '(CDeepStack)', False, 'from class_bind import AdjBinder\n'), ((168, 18, 168, 38), 'class_bind.AdjBinder', 'AdjBinder', ({(168, 28, 168, 37): 'CDeepBlob'}, {}), '(CDeepBlob)', False, 'from class_bind import AdjBinder\n'), ((170, 4, 170, 33), 'frame_blobs_yx.assign_adjacents', 'assign_adjacents', ({(170, 21, 170, 32): 'blob_binder'}, {}), '(blob_binder)', False, 'from frame_blobs_yx import assign_adjacents\n'), ((184, 9, 184, 16), 'collections.deque', 'deque', ({}, {}), '()', False, 'from collections import deque, defaultdict\n'), ((234, 17, 234, 29), 'utils.pairwise', 'pairwise', ({(234, 26, 234, 28): 'P_'}, {}), '(P_)', False, 'from utils import pairwise\n'), ((243, 14, 243, 21), 'collections.deque', 'deque', ({}, {}), '()', False, 'from collections import deque, defaultdict\n'), ((305, 18, 305, 25), 'collections.deque', 'deque', ({}, {}), '()', False, 'from collections import deque, defaultdict\n'), ((421, 11, 422, 56), 'numpy.pad', 'np.pad', (), '', True, 'import numpy as np\n'), ((100, 23, 100, 61), 'intra_comp_g.comp_r', 'comp_r', ({(100, 30, 100, 40): 'ext_dert__', (100, 42, 100, 45): 'fig', (100, 47, 100, 50): 'fcr', (100, 52, 100, 60): 'ext_mask'}, {}), '(ext_dert__, fig, fcr, ext_mask)', False, 'from intra_comp_g import comp_g, comp_r\n'), ((102, 23, 102, 51), 'intra_comp_g.comp_g', 'comp_g', ({(102, 30, 102, 40): 'ext_dert__', (102, 42, 102, 50): 'ext_mask'}, {}), '(ext_dert__, ext_mask)', False, 'from intra_comp_g import comp_g, comp_r\n'), ((151, 19, 151, 56), 'class_stream.BlobStreamer', 'BlobStreamer', ({(151, 32, 151, 41): 'CDeepBlob', (151, 43, 151, 49): 'crit__', (151, 51, 151, 55): 'mask'}, {}), '(CDeepBlob, crit__, mask)', False, 'from class_stream import BlobStreamer\n'), ((154, 19, 154, 56), 'class_stream.BlobStreamer', 'BlobStreamer', ({(154, 32, 154, 41): 'CDeepBlob', (154, 43, 154, 49): 'crit__', (154, 51, 154, 55): 'mask'}, {}), '(CDeepBlob, crit__, mask)', False, 'from class_stream import BlobStreamer\n'), ((157, 19, 157, 36), 'class_bind.AdjBinder', 'AdjBinder', ({(157, 29, 157, 35): 'CDeepP'}, {}), '(CDeepP)', False, 'from class_bind import AdjBinder\n'), ((384, 15, 384, 54), 'numpy.ones', 'np.ones', (), '', True, 'import numpy as np\n'), ((130, 26, 130, 84), 'itertools.zip_longest', 'zip_longest', (), '', False, 'from itertools import zip_longest\n')]