node_VectorMath2.py

# BEGIN GPL LICENSE BLOCK #####
#
#  This program is free software; you can redistribute it and/or
#  modify it under the terms of the GNU General Public License
#  as published by the Free Software Foundation; either version 2
#  of the License, or (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program; if not, write to the Free Software Foundation,
#  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# END GPL LICENSE BLOCK #####

import bpy
from node_s import *
from util import *
from mathutils import Vector, Matrix
from bpy.props import EnumProperty, StringProperty, BoolProperty

scalar_out = {
    "DOT":          (lambda u, v: u.dot(v), 2),
    "DISTANCE":     (lambda u, v: (u - v).length, 2),
    "ANGLE":        (lambda u, v: u.angle(v, 0), 2),
    "LEN":          (lambda u: u.length, 1)
}

vector_out = {
    "CROSS":        (lambda u, v: u.cross(v), 2),
    "ADD":          (lambda u, v: u + v, 2),
    "SUB":          (lambda u, v: u - v, 2),
    "REFLECT":      (lambda u, v: u.reflect(v), 2),
    "PROJECT":      (lambda u, v: u.project(v), 2),
    "SCALAR":       (lambda u, s: u * s, 2),
    "1/SCALAR":     (lambda u, s: u * (1/s), 2),
    "NORMALIZE":    (lambda u: u.normalized(), 1),
    "NEG":          (lambda u: -u, 1)
}


class VectorMath2Node(Node, SverchCustomTreeNode):

    ''' VectorMath 2 Node '''
    bl_idname = 'VectorMath2Node'
    bl_label = 'Vector Math 2'
    bl_icon = 'OUTLINER_OB_EMPTY'

    # vector math functions
    mode_items = [
        ("CROSS",       "Cross product",        "", 0),
        ("DOT",         "Dot product",          "", 1),
        ("ADD",         "Add",                  "", 2),
        ("SUB",         "Sub",                  "", 3),
        ("LEN",         "Length",               "", 4),
        ("DISTANCE",    "Distance",             "", 5),
        ("NORMALIZE",   "Normalize",            "", 6),
        ("NEG",         "Negate",               "", 7),
        ("ANGLE",       "Angle",                "", 12),
        ("PROJECT",     "Project",              "", 13),
        ("REFLECT",     "Reflect",              "", 14),
        ("SCALAR",      "Multiply Scalar",      "", 15),
        ("1/SCALAR",    "Multiply 1/Scalar",    "", 16)
    ]

    def mode_change(self, context):

        if not (self.items_ == self.current_op):
            self.label = self.items_
            self.update_outputs_and_inputs()
            self.current_op = self.items_
            updateNode(self, context)

    items_ = EnumProperty(
        items=mode_items,
        name="Function",
        description="Function choice",
        default="CROSS", update=mode_change)

    # matches default of CROSS product, defaults to False at init time.
    scalar_output_socket = BoolProperty()
    current_op = StringProperty(default="CROSS")

    def draw_buttons(self, context, layout):
        layout.prop(self, "items_", "Functions:")

    def init(self, context):
        self.inputs.new('VerticesSocket', "U", "u")
        self.inputs.new('VerticesSocket', "V", "v")
        self.outputs.new('VerticesSocket', "W", "W")

    def update_outputs_and_inputs(self):
        '''
        Reaches here only if new operation is different from current op
        '''
        inputs = self.inputs
        outputs = self.outputs
        new_op = self.items_
        current_op = self.current_op
        scalars = ["SCALAR", "1/SCALAR"]

        if (new_op in scalars) and (current_op in scalars):
            return  # it's OK nothing to change

        '''
        reaches this point, means it needs to rewire
        '''

        # check and adjust outputs and input size
        if new_op in scalar_out:
            self.scalar_output_socket = True
            nrInputs = scalar_out[new_op][1]
            if 'W' in outputs:
                outputs.remove(outputs['W'])
                outputs.new('StringsSocket', "out", "out")

        else:
            self.scalar_output_socket = False
            nrInputs = vector_out[new_op][1]
            if 'out' in outputs:
                outputs.remove(outputs['out'])
                outputs.new('VerticesSocket', "W", "W")

        '''
        this monster removes and adds sockets depending on what kind of switch
        between new_ops is made, some new_op changes require addition of sockets
        others require deletion or replacement.
        '''

        add_scalar_input = lambda: inputs.new('StringsSocket', "S", "s")
        add_vector_input = lambda: inputs.new('VerticesSocket', "V", "v")
        remove_last_input = lambda: inputs.remove(inputs[-1])

        if nrInputs < len(inputs):
            remove_last_input()

        elif nrInputs > len(inputs):
            if (new_op in scalars):
                add_scalar_input()
            else:
                add_vector_input()

        else:
            if nrInputs == 1:
                # is only ever a vector u
                return

            if new_op in scalars:
                remove_last_input()
                add_scalar_input()
            elif (current_op in scalars):
                remove_last_input()
                add_vector_input()

    def nothing_to_process(self):
        inputs = self.inputs
        outputs = self.outputs

        if not (('out' in outputs) or ('W' in outputs)):
            return True

        if not outputs[0].links:
            return True

    def update(self):
        inputs = self.inputs
        outputs = self.outputs
        operation = self.items_
        self.label = self.items_

        if self.nothing_to_process():
            return

        # this input is shared over both.
        vector1 = []
        if 'U' in inputs and inputs['U'].links:
            if isinstance(inputs['U'].links[0].from_socket, VerticesSocket):
                vector1 = SvGetSocketAnyType(self, inputs['U'])

        if not vector1:
            return

        # reaches here only if we have vector1
        u = vector1
        leve = levelsOflist(u)
        scalars = ["SCALAR", "1/SCALAR"]

        # vector-output
        if 'W' in outputs and outputs['W'].links:
            vector2, scalar1, result = [], [], []

            '''
            get input sockets content, depending on mode
            '''

            if 'V' in inputs and inputs['V'].links:
                if isinstance(inputs['V'].links[0].from_socket, VerticesSocket):
                    vector2 = SvGetSocketAnyType(self, inputs['V'])

            elif 'S' in inputs and inputs['S'].links:
                if isinstance(inputs['S'].links[0].from_socket, StringsSocket):
                    scalar1 = SvGetSocketAnyType(self, inputs['S'])

            '''
            figure out which functions and variables to use
            '''

            if len(inputs) == 1:

                recurse_fn = self.recurse_fx
                msg = 'one input only'

            elif len(inputs) == 2:

                # this means one of the necessary sockets is not connected
                if not (vector2 or scalar1):
                    return

                msg = "two inputs, "
                if operation in scalars:
                    b = scalar1
                    recurse_fn2 = self.recurse_fxy2
                    msg += "scalar"
                else:
                    b = vector2
                    recurse_fn2 = self.recurse_fxy
                    msg += 'non-scalar'

            else:
                return  # fail!

            '''
            At this point we've checked the input as much as we can,
            but Try just in case.
            '''
            func = vector_out[operation][0]
            try:
                if len(inputs) == 1:
                    result = recurse_fn(u, func, leve - 1)
                elif len(inputs) == 2:
                    result = recurse_fn2(u, b, func, leve - 1)
            except:
                print(self.name, msg, 'failed')

            SvSetSocketAnyType(self, 'W', result)

        # scalar-output
        if 'out' in outputs and outputs['out'].links:

            vector2, result = [], []
            func = scalar_out[operation][0]
            num_inputs = len(inputs)

            try:
                if num_inputs == 1:
                    result = self.recurse_fx(u, func, leve - 1)

                elif all([num_inputs == 2, ('V' in inputs), (inputs['V'].links)]):

                    if isinstance(inputs['V'].links[0].from_socket, VerticesSocket):
                        vector2 = SvGetSocketAnyType(self, inputs['V'])
                        result = self.recurse_fxy(u, vector2, func, leve - 1)
                    else:
                        print('socket connected to V is not a vertices socket')
                else:
                    return

            except:
                print('failed scalar out, {} inputs'.format(num_inputs))
                return

            if result:
                SvSetSocketAnyType(self, 'out', result)

    '''
    apply f to all values recursively
    - recurse_fxy taken from MathNode
    - fxy, fxy2 do full list matching by length
    is it best to keep these functions separate to avoid extra cycles per function call?
    they can be in very tight loops. Only really one way to find out!

    '''

    # vector -> scalar | vector
    def recurse_fx(self, l, f, leve):
        if not leve:
            w = f(Vector(l))
            if self.scalar_output_socket:
                return w
            else:
                return w.to_tuple()
        else:
            t = []
            for i in l:
                t.append(self.recurse_fx(i, f, leve - 1))
        return t

    # vector, vector -> scalar | vector
    def recurse_fxy(self, l1, l2, f, leve):
        if not leve:
            w = f(Vector(l1), Vector(l2))
            if self.scalar_output_socket:
                return w
            else:
                return w.to_tuple()
        else:
            max_obj = max(len(l1), len(l2))
            fullList(l1, max_obj)
            fullList(l2, max_obj)
            res = []
            for i in range(len(l1)):
                res.append(self.recurse_fxy(l1[i], l2[i], f, leve - 1))
            return res

    # vector, scalar -> vector
    def recurse_fxy2(self, l1, l2, f, leve):
        if not leve:
            w = f(Vector(l1), l2)
            return w.to_tuple()
        else:
            max_obj = max(len(l1), len(l2))
            fullList(l1, max_obj)
            fullList(l2, max_obj)
            res = []
            for i in range(len(l1)):
                res.append(self.recurse_fxy2(l1[i], l2[i], f, leve - 1))
            return res

    def update_socket(self, context):
        self.update()


def register():
    bpy.utils.register_class(VectorMath2Node)


def unregister():
    bpy.utils.unregister_class(VectorMath2Node)

if __name__ == "__main__":
    register()