bTree

# bSearch v1
# some kind of ocTree thing for maya. Almost working, currently just subdivides space
# 2-12-2017

from pymel.core import *
import random as r
import maya.api.OpenMaya as om
import maya.cmds as mc
import time
start = time.clock()
class defAttr:
    subDepth = 2
    subMaxCount = 20
    dCheck = 20
    showGrow = 1

    saveAll = []
    toParent = []
    showObjs = []
    currentSet = []
    lastSet = []
    setSave = []
    searchList = []
class searchObj():
    def __init__(self,obj):
        self.sPos = obj.getTranslation()
        self.isFound = 0
        self.sObj = []
def getDist(p1,p2):
    dist = math.sqrt((p1[0] - p2[0])**2 + (p1[1] - p2[1])**2 + (p1[2] - p2[2])**2)
    return dis
def grpr(name):
    g = ls(name)
    if g:
        delete(g)
    g = group(n=name,em=1)
    return g
class subBB:
    def __init__(self,bbb):
        self.bb = bbb
        self.bbChildren = []
        self.bbCenter = (bbb[0]+bbb[3])/2, (bbb[1]+bbb[4])/2, (bbb[2]+bbb[5])/2
        self.bbObj = []
        self.bbActive = 0
        self.bbIsInObjs = []
        self.bbSearchObjs = []
        self.pap = []
    def doSubBB(self):#xmin0 ymin1 zmin2 xmax3 ymax4 zmax5
        bb = self.bb
        newBb1 = subBB([bb[0], bb[1], bb[2], (bb[3]+bb[0])/2, (bb[4]+bb[1])/2, (bb[5]+bb[2])/2])
        newBb2 = subBB([(bb[0]+bb[3])/2, bb[1], bb[2], bb[3], (bb[4]+bb[1])/2, (bb[5]+bb[2])/2])
        newBb3 = subBB([(bb[0]+bb[3])/2, bb[1], (bb[2]+bb[5])/2, bb[3], (bb[4]+bb[1])/2, bb[5]])
        newBb4 = subBB([bb[0], bb[1], (bb[2]+bb[5])/2, (bb[3]+bb[0])/2, (bb[4]+bb[1])/2, bb[5]])
        newBbUp1 = subBB([bb[0], (bb[1]+bb[4])/2, bb[2], (bb[3]+bb[0])/2, bb[4], (bb[5]+bb[2])/2])
        newBbUp2 = subBB([(bb[0]+bb[3])/2, (bb[1]+bb[4])/2, bb[2], bb[3], bb[4], (bb[5]+bb[2])/2])
        newBbUp3 = subBB([(bb[0]+bb[3])/2, (bb[1]+bb[4])/2, (bb[2]+bb[5])/2, bb[3], bb[4], bb[5]])
        newBbUp4 = subBB([bb[0], (bb[1]+bb[4])/2, (bb[2]+bb[5])/2, (bb[3]+bb[0])/2, bb[4], bb[5]])
        newBBList = [newBb1,newBb2,newBb3,newBb4,newBbUp1,newBbUp2,newBbUp3,newBbUp4]
        addList = []
        for new in newBBList:
            new.bbSearchObjs = self.bbSearchObjs
            new.isIn()
            if new.bbActive:
                addList.append(new)
                defAttr.saveAll.append(new)
        self.bbChildren = addList

    def doGrow(self):
        newBB = self.bb
        i = 0
        c = polyCube(w=0,h=0,d=0, n = "showCube_"+str(i))
        mc.move(newBB[0]+.5,0,0, str(c[0])+".f[5]", r=1,wd=1,ws=1)
        mc.move(0,newBB[1]+.5,0, str(c[0])+".f[3]", r=1,wd=1,ws=1)
        mc.move(0,0,newBB[2]+.5, str(c[0])+".f[2]", r=1,wd=1,ws=1)
        mc.move(newBB[3]-.5,0,0, str(c[0])+".f[4]",  r=1,wd=1,ws=1)
        mc.move(0,newBB[4]-.5,0, str(c[0])+".f[1]",  r=1,wd=1,ws=1)
        mc.move(0,0,newBB[5]-.5, str(c[0])+".f[0]",  r=1,wd=1,ws=1)
        defAttr.showObjs.append(c[0])
        self.bbObj = [c]
        i += 1
    def isIn(self):#(bb,obj):#xmin ymin zmin xmax ymax zmax.
        isInList = []
        bb = self.bb
        for o in self.bbSearchObjs:
            if not o.isFound:
                add = 0
                objPos = o.sPos
                if objPos[0] <= bb[0] or objPos[0] >= bb[3]:
                    add = 0
                elif objPos[1] <= bb[1] or objPos[1] >= bb[4]:
                    add = 0
                elif objPos[2] <= bb[2] or objPos[2] >= bb[5]:
                    add = 0
                else:
                    add=1
                if add:
                    self.bbIsInObjs.append(o)
                    if len(self.bbSearchObjs) <= defAttr.subMaxCount:
                        self.bbActive = 0
                    else:
                        self.bbActive = 1
                    isInList.append(o)
        self.bbSearchObjs = isInList

def isInSimple(finder,pp):#(bb,obj):#xmin ymin zmin xmax ymax zmax.
    add = 0
    bb = finder.bb
    objPos = pp

    isIn = []
    if objPos[0] <= bb[0] or objPos[0] >= bb[3]:
        add = 0
    elif objPos[1] <= bb[1] or objPos[1] >= bb[4]:
        add = 0
    elif objPos[2] <= bb[2] or objPos[2] >= bb[5]:
        add = 0
    else:
        add=1

    if add:
        isIn.append(finder)
        finder.bbActive = 1
    else:
        finder.bbActive = 0

def getClose(closeToObj):

    finder = defAttr.saveAll[0]
    p2 = om.MVector(closeToObj.getTranslation())
    p2R = closeToObj.getScale()[0]
    select(cl=1)
    toSl = []
    checkList = finder.bbChildren
    for c in checkList:
        isInSimple(c, p2)
        if not c.bbActive:
            newDir = (om.MVector(p2) + om.MVector(c.bbCenter)*.5) - om.MVector(c.bbCenter)
            pX = om.MVector(p2) - om.MVector.normalize(newDir)*p2R
            isInSimple(c, pX)
            if c.bbActive:
                toSl.append(c.bbObj[0])
        else:
            toSl.append(c.bbObj[0])
    return toSl

subGrpr = grpr("subGrpr")
#sl = selected()
leafGrp = ls("ThingGrp")[0]
leafGrpC = leafGrp.listRelatives(c=1)
newLeafGrpC = []
for leaf in leafGrpC:
    newL = searchObj(leaf)
    newL.sObj = leaf
    newLeafGrpC.append(newL)

bbb = xform(leafGrp, bbi=1, q=1)
scSet = subBB(bbb)
scSet.bbSearchObjs = newLeafGrpC
defAttr.currentSet = [scSet]
defAttr.saveAll = [scSet]
progWin = progressWindow(isInterruptable=1, t= "Simulating...", min = 0, max = defAttr.subDepth+1)
ii = 0
for i in range(0,defAttr.subDepth+1):
    progressWindow(progWin, edit=True, pr = i)
    toSelect = []
    defAttr.lastSet = []
    for cSet in defAttr.currentSet:# list of bounding boxs
        if progressWindow(progWin, query=1, isCancelled=1 ):
            print "killed at", i+1
            break
        cSet.doSubBB()
        defAttr.setSave.append(cSet.bbChildren)
        defAttr.lastSet.append(cSet)
        if defAttr.showGrow:
            cSet.doGrow()
            toSelect.append(cSet.bbObj[0])
            select(cl=1)
            refresh(cw=1)
    listSaveFormat = []
    for ss in defAttr.setSave:
        for sss in ss:
            listSaveFormat.append(sss)
    for cc in defAttr.searchList:
        cc.isFound = 0
    defAttr.currentSet = listSaveFormat
    defAttr.setSave = []
    tempLast = defAttr.lastSet
    ii += 1
parent(defAttr.showObjs, subGrpr)
progressWindow(progWin, endProgress=1)
select(cl=1)


closeToObj = ls("close_obj")[0]
closeList = getClose(closeToObj)
#toS = []
#print defAttr.saveAll[0].bbChildren[0].bbCenter
print closeList
select(closeList)


"""
for o in defAttr.saveAll:
    if o.bbChildren:
        toS.append(o.bbChildren)
select(toS)
"""

elapsed = (time.clock() - start)
elapsed = elapsed
print "\n","Run time total:", elapsed