from maya import cmds
import maya.api.OpenMaya as om
from avalon.pipeline import AVALON_CONTAINER_ID

import avalon.maya.pipeline

# Current implementation
ls = avalon.maya.pipeline._ls


def ls1():
    """First query to only object sets and attribute .id

    Then use API to get string value of the .id
    """

    nodes = cmds.ls("*.id", exactType="objectSet", o=True, recursive=True)

    ids = {AVALON_CONTAINER_ID,
           # Backwards compatibility
           "pyblish.mindbender.container"}

    if not nodes:
        return

    sel = om.MSelectionList()
    for node in nodes:
        sel.add(node)

    fn_dep = om.MFnDependencyNode()
    for i, node in enumerate(nodes):
        dep_node = sel.getDependNode(i)
        fn_dep.setObject(dep_node)
        if fn_dep.findPlug("id", True).asString() in ids:
            yield node

def ls2():
    """Iterate with API over all set types"""

    def _maya_iterate(iterator):
        """Helper to just iterate a maya iterator"""
        while not iterator.isDone():
            yield iterator.thisNode()
            iterator.next()

    ids = {AVALON_CONTAINER_ID,
           # Backwards compatibility
           "pyblish.mindbender.container"}

    # Getting the string attribute values is much faster through the Maya API
    # so the following might seem convoluted but it's much faster.
    fn_dep = om.MFnDependencyNode()
    iterator = om.MItDependencyNodes(om.MFn.kSet)
    for mobject in _maya_iterate(iterator):
        if mobject.apiTypeStr != "kSet":
            # Only match by exact type
            continue

        fn_dep.setObject(mobject)
        try:
            plug = fn_dep.findPlug("id", True)
        except RuntimeError:
            continue

        value = plug.asString()
        if value in ids:
            yield fn_dep.name()

def ls3():

    nodes = cmds.ls(exactType="objectSet", recursive=True)

    ids = {AVALON_CONTAINER_ID,
           # Backwards compatibility
           "pyblish.mindbender.container"}

    if not nodes:
        return

    sel = om.MSelectionList()
    for node in nodes:
        sel.add(node)

    fn_dep = om.MFnDependencyNode()
    for i, node in enumerate(nodes):
        mobject = sel.getDependNode(i)
        fn_dep.setObject(mobject)
        try:
            plug = fn_dep.findPlug("id", True)
        except RuntimeError:
            continue

        value = plug.asString()
        if value in ids:
            yield fn_dep.name()



s = time.time()
result = list(ls())
e = time.time()
print "Current:", e-s

s = time.time()
result1 = list(ls1())
e = time.time()
print "New 1:", e-s

s = time.time()
result2 = list(ls2())
e = time.time()
print "New 2:", e-s

s = time.time()
result3 = list(ls3())
e = time.time()
print "New 3:", e-s

# Sort all to ensure equality check is consistent
result.sort()
result1.sort()
result2.sort()
result3.sort()

assert result1 == result, "ls1() does not match ls()"
assert result2 == result, "ls2() does not match ls()"
assert result3 == result, "ls3() does not match ls()"