''' mbinary
#########################################################################
# File : bTree.py
# Author: mbinary
# Mail: zhuheqin1@gmail.com
# Blog: https://mbinary.coding.me
# Github: https://github.com/mbinary
# Created Time: 2018-08-29  12:49
# Description:
#########################################################################
'''

class node:
    def __init__(self,keys=None,isLeaf = True,children=None):
        if keys is None:keys=[]
        if children is None: children =[]
        self.keys = keys
        self.isLeaf =  isLeaf
        self.children = []
    def __getitem__(self,i):
        return self.keys[i]
    def __delitem__(self,i):
        del self.keys[i]
    def __setitem__(self,i,k):
        self.keys[i] = k
    def __len__(self):
        return len(self.keys)
    def __repr__(self):
        return str(self.keys)
    def __str__(self):
        children = ','.join([str(nd.keys) for nd in self.children])
        return f'keys:     {self.keys}\nchildren: {children}\nisLeaf:   {self.isLeaf}'
    def getChd(self,i):
        return self.children[i]
    def delChd(self,i):
        del self.children[i]
    def setChd(self,i,chd):
        self.children[i] = chd
    def getChildren(self,begin=0,end=None):
        if end is None:return self.children[begin:]
        return self.children[begin:end]
    def findKey(self,key):
        for i,k in enumerate(self.keys):
            if k>=key:
                return i
        return len(self)
    def update(self,keys=None,isLeaf=None,children=None):
        if keys is not None:self.keys = keys
        if children is not None:self.children = children
        if isLeaf is not None: self.isLeaf = isLeaf
    def insert(self,i,key=None,nd=None):
        if key is not None:self.keys.insert(i,key)
        if not self.isLeaf and nd is not None: self.children.insert(i,nd)
    def isLeafNode(self):return self.isLeaf
    def split(self,prt,t):
        # form new two nodes
        k = self[t-1]
        nd1 = node()
        nd2 = node()
        nd1.keys,nd2.keys = self[:t-1], self[t:] # note that t is 1 bigger than  key index
        nd1.isLeaf = nd2.isLeaf = self.isLeaf
        if not self.isLeaf: 
            # note that  children index is one bigger than key index, and all children included
            nd1.children, nd2.children = self.children[0:t], self.children[t:] 
        # connect them to parent
        idx = prt.findKey(k)
        if prt.children !=[]: prt.children.remove(self) # remove the original node
        prt.insert(idx,k,nd2)
        prt.insert(idx,nd = nd1)
        return prt


class bTree:
    def __init__(self,degree=2):
        self.root = node()
        self.degree=degree
        self.nodeNum = 1
        self.keyNum = 0
    def search(self,key,withpath=False):
        nd = self.root
        fathers = []
        while True:
            i = nd.findKey(key)
            if i==len(nd): fathers.append((nd,i-1,i))
            else: fathers.append((nd,i,i))
            if i<len(nd) and nd[i]==key:
                if withpath:return nd,i,fathers
                else:return nd,i
            if nd.isLeafNode(): 
                if withpath:return None,None,None
                else:return None,None
            nd = nd.getChd(i)
    def insert(self,key):
        if len(self.root)== self.degree*2-1:
            self.root = self.root.split(node(isLeaf=False),self.degree)
            self.nodeNum +=2
        nd = self.root
        while True:
            idx = nd.findKey(key)
            if idx<len(nd) and nd[idx] == key:return
            if nd.isLeafNode():
                nd.insert(idx,key)
                self.keyNum+=1
                return
            else:
                chd = nd.getChd(idx)
                if len(chd)== self.degree*2-1: #ensure its keys won't excess when its chd split and u
                    nd = chd.split(nd,self.degree)
                    self.nodeNum +=1
                else:
                    nd = chd
    def delete(self,key):#to do
        '''search the key, delete it , and form down to up to rebalance it '''
        nd,idx ,fathers= self.search(key,withpath=True)
        if nd is None : return
        del nd[idx]
        self.keyNum-=1
        if not nd.isLeafNode():
            chd = nd.getChd(idx) # find the predecessor key
            while not  chd.isLeafNode():
                fathers.append((chd,len(chd)-1,len(chd)))
                chd = chd.getChd(-1)
            fathers.append((chd,len(chd)-1,len(chd)))
            nd.insert(idx,chd[-1])
            del chd[-1]
        if len(fathers)>1:self.rebalance(fathers)
    def rebalance(self,fathers):
        nd,keyIdx,chdIdx = fathers.pop()
        while len(nd)<self.degree-1: # rebalance tree from down to up
            prt,keyIdx,chdIdx = fathers[-1]
            lbro = [] if chdIdx==0 else prt.getChd(chdIdx-1)
            rbro = [] if chdIdx==len(prt) else prt.getChd(chdIdx+1)
            if len(lbro)<self.degree and len(rbro)<self.degree:  # merge two deficient nodes
                beforeNode,afterNode = None,None
                if lbro ==[]:
                    keyIdx = chdIdx
                    beforeNode,afterNode = nd,rbro
                else:
                    beforeNode,afterNode = lbro,nd
                    keyIdx = chdIdx-1      # important, when choosing 
                keys = beforeNode[:]+[prt[keyIdx]]+afterNode[:]
                children = beforeNode.getChildren() + afterNode.getChildren()
                isLeaf = beforeNode.isLeafNode()
                prt.delChd(keyIdx+1)
                del prt[keyIdx]
                nd.update(keys,isLeaf,children)
                prt.children[keyIdx]=nd
                self.nodeNum -=1
            elif len(lbro)>=self.degree:  # rotate  when only one sibling is deficient
                keyIdx = chdIdx-1
                nd.insert(0,prt[keyIdx])    # rotate keys
                prt[keyIdx] =  lbro[-1]
                del lbro[-1]
                if not nd.isLeafNode():     # if not leaf, move children
                    nd.insert(0,nd=lbro.getChd(-1))
                    lbro.delChd(-1)
            else:
                keyIdx = chdIdx
                nd.insert(len(nd),prt[keyIdx])    # rotate keys
                prt[keyIdx] =  rbro[0]
                del rbro[0]
                if not nd.isLeafNode():     # if not leaf, move children
                    #note that insert(-1,ele) will make the ele be the last second one
                    nd.insert(len(nd),nd=rbro.getChd(0))
                    rbro.delChd(0)
            if len(fathers)==1:
                if len(self.root)==0: 
                    self.root = nd
                    self.nodeNum -=1
                break
            nd,i,j = fathers.pop()
    def __str__(self):
        head= '\n'+'-'*30+'B  Tree'+'-'*30
        tail= '-'*30+'the end'+'-'*30+'\n'
        lst = [[head],[f'node num: {self.nodeNum},  key num: {self.keyNum}']]
        cur = []
        ndNum =0
        ndTotal= 1
        que = [self.root]
        while que!=[]:
            nd = que.pop(0)
            cur.append(repr(nd))
            ndNum+=1
            que+=nd.getChildren()
            if ndNum==ndTotal:
                lst.append(cur)
                cur = []
                ndNum = 0
                ndTotal =len(que) 
        lst.append([tail])
        lst = [','.join(li) for li in lst]
        return '\n'.join(lst)
    def __iter__(self,nd = None):
        if nd is None: nd = self.root
        que = [nd]
        while que !=[]:
            nd = que.pop(0)
            yield nd
            if nd.isLeafNode():continue
            for i in range(len(nd)+1):
                que.append(nd.getChd(i))


if __name__ =='__main__':
    bt = bTree()
    from random import shuffle,sample
    n = 20
    lst = [i for i in range(n)]
    shuffle(lst)
    test= sample(lst,len(lst)//4)
    print(f'building b-tree with  {lst}')
    for i in lst:
        bt.insert(i)
        #print(f'inserting {i})
        #print(bt)
    print(bt)
    print(f'serching {test}')
    for i in test:
        nd,idx = bt.search(i)
        print(f'node: {repr(nd)}[{idx}]== {i}')
    for i in test:
        print(f'deleting {i}')
        bt.delete(i)
        print(bt)