Split Flap Display

splitflap.py

from __future__ import print_function, division
import pymel.core as pm
from pymel.core.runtime import nClothCreate, nClothMakeCollide
from contextlib import contextmanager
import re
import time
from random import choice


@contextmanager
def selection(nodes):
    old_selection = pm.selected()
    try:
        pm.select(nodes, replace=True)
        yield
    finally:
        pm.select(old_selection, replace=True)


@contextmanager
def undo_chunk(auto_undo=False, exc_callback=None):

    try:
        pm.undoInfo(openChunk=True)
        yield
    except Exception as e:
        pm.undoInfo(closeChunk=True)
        if auto_undo:
            pm.undo()
        if exc_callback:
            exc_callback(e)
        raise
    else:
        pm.undoInfo(closeChunk=True)
        if auto_undo:
            pm.undo()


def set_uvs(mesh, uvs, uvSet=None):
    uvSet = uvSet or 'map1'
    mesh.setUVs(*zip(*uvs), uvSet=uvSet)


def get_uvs(mesh):
    return zip(*mesh.getUVs())


def get_uvs_in_range(uvs, u_min, v_min, u_max, v_max):
    uvids = []
    for i, (u, v) in enumerate(uvs):
        if u_min < u < u_max and v_min < v < v_max:
            uvids.append(i)
    return uvids


def get_row_col(index, max_index, num_columns):
    '''
    Get the row and coloumn for the specified index.

    :param index: Index of item
    :param max_index: Maximum number of items in loop
    :param num_colums: Number of columns
    '''

    if max_index:
        if index >= max_index:
            index -= max_index

    row = index/(num_columns)
    col = index%(num_columns)
    return floor(row), floor(col)


def pack_uvs(uvs, uvids, i, rows, columns):
    r, c = get_row_col(i, None, columns)
    sx = 1 / columns
    sy = 1 / rows

    for uvid in uvids:
        u, v = uvs[uvid]
        u = u * sx + c * sx
        v = v * sy + (1 - sy - r * sy)
        uvs[uvid] = (u, v)


def index_to_udim(i, max_i, columns=10):
    r, c = get_row_col(i, max_i, columns)
    return c, r


def shift_uvs(uvs, uvids, u_shift, v_shift):

    for uvid in uvids:
        u, v = uvs[uvid]
        u += u_shift
        v += v_shift
        uvs[uvid] = (u, v)


def create_cloth_flap(flap, subdivisions_width=1, subdivisions_height=1):
    bounds = flap.boundingBox()
    plane, plane_shape = pm.polyPlane(
        width=bounds.width(),
        height=bounds.height(),
        sx=subdivisions_width,
        sy=subdivisions_height,
        axis=[0, 0, 1],
        cuv=1,
    )
    plane.setTranslation(bounds.center())
    pm.makeIdentity(plane, apply=True, t=True, r=True, s=True, n=False)
    plane.setPivots([0, 0, 0])
    return plane


def create_flaps(num_flaps, base_flaps, layout_index, rows, columns):

    # Name and group geometry
    r, c = get_row_col(layout_index, None, columns)
    flaps = [choice(base_flaps).duplicate(rc=True)[0]
             for i in xrange(num_flaps)]

    # Pack UVS
    uvs = get_uvs(flaps[0].getShape(noIntermediate=True))
    top_uvids = get_uvs_in_range(uvs, 0, 0.5, 1, 1)
    bottom_uvids = get_uvs_in_range(uvs, 0, 0, 1, 0.5)
    pack_uvs(
        uvs,
        top_uvids + bottom_uvids,
        layout_index,
        rows,
        columns
    )

    for i, flap in enumerate(flaps):
        mesh = flap.getShape(noIntermediate=True)

        # Shift UVS according to UDIM
        mesh_uvs = list(uvs)
        shift_uvs(mesh_uvs, top_uvids, *index_to_udim(i, num_flaps))
        shift_uvs(mesh_uvs, bottom_uvids, *index_to_udim(i + 1, num_flaps))
        set_uvs(mesh, mesh_uvs)

    return flaps


def radial_arrangement(transforms, radius):

    rotate_step = 360 / len(transforms)
    radians = math.pi / 180

    for i, t in enumerate(transforms):
        rx = rotate_step * i
        theta = rx * radians
        translation = [0, math.cos(theta) * radius, -math.sin(theta) * radius]
        pm.xform(t, rotation=[-rx, 0, 0], translation=translation)


def create_wrap_deformer(influence, deformed, **kwargs):

    pm.select(deformed, replace=True)

    kwargs.setdefault('weightThreshold', 0.0)
    kwargs.setdefault('maxDistance', 1.0)
    kwargs.setdefault('exclusiveBind', False)
    kwargs.setdefault('autoWeightThreshold', True)
    kwargs.setdefault('falloffMode', 0)

    wrap = pm.deformer(type='wrap')[0]
    for k, v in kwargs.iteritems():
        wrap.attr(k).set(v)

    if not influence.hasAttr('dropoff'):
        influence.addAttr('dropoff', sn='dr', dv=4, min=0, max=20, k=True)
    if not influence.hasAttr('smoothness'):
        influence.addAttr('smoothness', sn='smt', dv=0, min=0, k=True)
    if not influence.hasAttr('inflType'):
        influence.addAttr('inflType', sn='ift', at='short', dv=2, min=1, max=2)

    influence.dropoff.connect(wrap.dropoff[0])
    influence.smoothness.connect(wrap.smoothness[0])
    influence.inflType.connect(wrap.inflType[0])

    influence_shape = influence.getShape(noIntermediate=True)
    influence_shape.worldMesh.connect(wrap.driverPoints[0])

    base = influence.duplicate(name=influence + 'Shape', rc=True)[0]
    base_shape = base.getShape(noIntermediate=True)
    base.hide()
    base_shape.worldMesh.connect(wrap.basePoints[0])

    deformed.worldMatrix.connect(wrap.geomMatrix)

    return wrap, base

def create_collider(cloth, radius):
    tx = cloth.boundingBox().width() * 0.5
    ty = cloth.boundingBox().height() * 0.5 - radius
    tz = radius * 0.9
    cubea, cubea_shape = pm.polyCube(width=0.2, height=0.1, depth=0.01)
    cubeb, cubeb_shape = pm.polyCube(width=0.2, height=0.1, depth=0.01)
    cubea.setTranslation([tx, ty, tz])
    cubeb.setTranslation([-tx, ty, tz])
    collider = pm.polyUnite(
        [cubea, cubeb],
        ch=False,
        mergeUVSets=True,
        name=cloth.replace('geo', 'collider_geo')
    )[0]
    return collider


def create_split_flap(base_flaps, num_images, rows, columns, radius,
                      layout_index=0):

    flaps = create_flaps(num_images, base_flaps, layout_index, rows, columns)
    cloth_flaps = [create_cloth_flap(flaps[0])]
    cloth_flaps.extend([cloth_flaps[0].duplicate(rc=True)[0]
                        for i in xrange(num_images - 1)])
    radial_arrangement(flaps, radius)
    radial_arrangement(cloth_flaps, radius)

    r, c = get_row_col(layout_index, None, columns)
    rowcol = '{:02d}{:02d}'.format(int(r), int(c))
    cloth_name = 'cloth_flap_{}'.format(rowcol)
    flaps_name = 'flaps_{}'.format(rowcol)

    cloth = pm.polyUnite(
        cloth_flaps,
        ch=False,
        mergeUVSets=True,
        name=cloth_name + '_geo',
    )[0]
    flaps = pm.polyUnite(
        flaps,
        ch=False,
        mergeUVSets=True,
        name=flaps_name + '_geo',
    )[0]
    pm.hide(cloth)

    # Create colliders
    collider = create_collider(flaps, radius)

    flaps_grp = pm.group([flaps, collider],
                         name='flaps_{}_geo_grp'.format(rowcol))
    rotate_grp = pm.group(cloth, name='rotate_{}_grp'.format(rowcol))
    split_flap = pm.group([rotate_grp, flaps_grp], name=flaps_name + '_grp')
    split_flap.addAttr('split_flap', at='bool', dv=True)
    split_flap.addAttr('layout_index', at='long', dv=layout_index)
    split_flap.addAttr('layout_row', at='long', dv=r)
    split_flap.addAttr('layout_column', at='long', dv=c)
    split_flap.addAttr('flaps', at='message')
    split_flap.addAttr('cloth', at='message')
    split_flap.addAttr('collider', at='message')
    flaps.message.connect(split_flap.flaps)
    cloth.message.connect(split_flap.cloth)
    collider.message.connect(split_flap.collider)

    return SplitFlap(split_flap)


def create_split_flap_wall(split_flap, rows, columns, padding=0.2):
    bounds = split_flap.pynode.boundingBox()
    x_step = bounds.width() + padding
    y_step = bounds.height() + padding
    u_step = 1 / columns
    v_step = 1 / rows
    uvs = get_uvs(split_flap.flaps)
    uvids = get_uvs_in_range(uvs, 0, 0, 9999, 9999)

    split_flaps = []
    i = 0
    for r in xrange(rows):
        for c in xrange(columns):
            if r == c == 0:
                continue
            index_name = '{:02d}{:02d}'.format(r, c)
            new_flap = split_flap.pynode.duplicate(
                name='flaps_{}_grp'.format(index_name),
                un=True,
                rc=False)[0]

            # Translate group
            new_flap.setTranslation([c * x_step, r * -y_step, 0])
            new_flap.layout_index.set(i)
            new_flap.layout_row.set(r)
            new_flap.layout_column.set(c)

            # Shift uvs
            new_split_flap = SplitFlap(new_flap)
            flaps = new_split_flap.flaps
            mesh_uvs = list(uvs)
            shift_uvs(mesh_uvs, uvids, c * u_step, r * -v_step)
            set_uvs(flaps, mesh_uvs)

            # Rename hierarchy
            hierarchy = pm.ls(new_flap, dag=True)
            for node in hierarchy:
                old_name = str(node)
                new_name = re.sub(r'\d+', index_name, old_name)
                if old_name != new_name:
                    node.rename(new_name)

            split_flaps.append(new_split_flap)
            i += 1

    grp = pm.group([f.pynode for f in split_flaps], name='split_flap_wall')
    pm.parent(split_flap.pynode, grp)
    return SplitFlapWall(grp)


class SplitFlapWall(object):

    def __init__(self, pynode):
        self.pynode = pynode
        self._split_flaps = None

    @property
    def split_flaps(self):
        if not self._split_flaps:
            self._split_flaps = []
            for node in pm.ls(self.pynode, dag=True):
                if node.hasAttr('split_flap'):
                    self._split_flaps.append(SplitFlap(node))
        return self._split_flaps

    def make_dynamic(self):
        needs_ncloth = []
        colliders = []

        for split_flap in self.split_flaps:
            if not split_flap.is_dynamic:
                needs_ncloth.append(split_flap.cloth)
                colliders.append(split_flap.collider)

        if needs_ncloth:
            with selection(needs_ncloth):
                nClothCreate()
                ncloth_shapes = pm.selected()
                for ncloth in ncloth_shapes:
                    ncloth.inputMeshAttract.set(1)
                    ncloth.inputAttractMethod.set(1)
                    ncloth.inputAttractDamp.set(0)
                    ncloth.thickness.set(0.005)
                    ncloth.pointMass.set(10)
                    ncloth.drag.set(0.5)
                ncloth_transforms = [shape.getParent()
                                     for shape in ncloth_shapes]

            with selection(colliders):
                nClothMakeCollide()
                collider_shapes = pm.selected()
                collider_transforms = [shape.getParent()
                                       for shape in collider_shapes]

            ramp = pm.createNode('ramp')
            ramp.colorEntryList[0].color.set(1, 1, 1)
            ramp.colorEntryList[1].position.set(0.1)
            ramp.colorEntryList[1].color.set(0, 0, 0)
            ramp.interpolation.set(0)

            for shape in ncloth_shapes:
                ramp.outAlpha.connect(shape.inputAttractMap)

        wrap_nodes = []
        for split_flap in self.split_flaps:
            _, base = create_wrap_deformer(split_flap.cloth, split_flap.flaps)
            wrap_nodes.append(base)

        pm.group(ncloth_transforms + collider_transforms,
                 name='dynamics_grp',
                 parent=self.pynode)


class SplitFlap(object):

    def __init__(self, pynode):
        self.pynode = pynode
        self._flaps = None
        self._cloth = None
        self._collider = None

    @property
    def layout_index(self):
        return self.pynode.layout_index

    @property
    def layout_row(self):
        return self.pynode.layout_row

    @property
    def layout_column(self):
        return self.pynode.layout_column

    @property
    def flaps(self):
        if not self._flaps:
            self._flaps = self.pynode.flaps.inputs()[0]
        return self._flaps

    @property
    def cloth(self):
        if not self._cloth:
            self._cloth = self.pynode.cloth.inputs()[0]
        return self._cloth

    @property
    def collider(self):
        if not self._collider:
            self._collider = self.pynode.collider.inputs()[0]
        return self._collider

    @property
    def is_dynamic(self):
        return bool(self.cloth.history(type='nCloth'))

    def ncloth_shape(self):
        cloth_shape = self.cloth.getShape(noIntermediate=True)
        return cloth_shape.inMesh.inputs(type='nCloth')[0]


if __name__ == '__main__':
    with undo_chunk():
        split_flap = create_split_flap(
            base_flaps=pm.selected(),
            num_images=16,
            rows=5,
            columns=10,
            radius=0.2)

    with undo_chunk():
        split_flap_wall = create_split_flap_wall(
            split_flap,
            rows=5,
            columns=10,
            padding=0.2)

        # split_flap_wall.make_dynamic()