utils.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 #####

print("version 3 imported")

import bpy
import time
import copy

from mathutils import *
from math import pi, sin, degrees, radians, atan2, copysign, cos, acos
from math import floor, ceil
from random import random, uniform, seed, choice, getstate, setstate, randint
from bpy.props import *
from collections import deque, OrderedDict

tau = 2 * pi

# Initialise the split error and axis vectors
splitError = 0.0
zAxis = Vector((0, 0, 1))
yAxis = Vector((0, 1, 0))
xAxis = Vector((1, 0, 0))

# This class will contain a part of the tree which needs to be extended and the required tree parameters
class stemSpline:
    def __init__(self, spline, curvature, curvatureV, attractUp, segments, maxSegs, segLength, childStems, stemRadStart, stemRadEnd, splineNum, ofst, pquat):
        self.spline = spline
        self.p = spline.bezier_points[-1]
        self.curv = curvature
        self.curvV = curvatureV
        self.vertAtt = attractUp
        self.seg = segments
        self.segMax = maxSegs
        self.segL = segLength
        self.children = childStems
        self.radS = stemRadStart
        self.radE = stemRadEnd
        self.splN = splineNum
        self.offsetLen = ofst
        self.patentQuat = pquat
        self.curvSignx = 1
        self.curvSigny = 1

    # This method determines the quaternion of the end of the spline
    def quat(self):
        if len(self.spline.bezier_points) == 1:
            return ((self.spline.bezier_points[-1].handle_right - self.spline.bezier_points[-1].co).normalized()).to_track_quat('Z', 'Y')
        else:
            return ((self.spline.bezier_points[-1].co - self.spline.bezier_points[-2].co).normalized()).to_track_quat('Z', 'Y')
    # Determine the declination
    def dec(self):
        tempVec = zAxis.copy()
        tempVec.rotate(self.quat())
        return zAxis.angle(tempVec)
    # Update the end of the spline and increment the segment count
    def updateEnd(self):
        self.p = self.spline.bezier_points[-1]
        self.seg += 1

# This class contains the data for a point where a new branch will sprout
class childPoint:
    def __init__(self, coords, quat, radiusPar, offset, sOfst, lengthPar, parBone):
        self.co = coords
        self.quat = quat
        self.radiusPar = radiusPar
        self.offset = offset
        self.stemOffset = sOfst
        self.lengthPar = lengthPar
        self.parBone = parBone


# This function calculates the shape ratio as defined in the paper
def shapeRatio(shape, ratio, pruneWidthPeak=0.0, prunePowerHigh=0.0, prunePowerLow=0.0, custom=None):
    if shape == 0:
        return 0.05 + 0.95*ratio #0.2 + 0.8*ratio
    elif shape == 1:
        return 0.2 + 0.8*sin(pi*ratio)
    elif shape == 2:
        return 0.2 + 0.8*sin(0.5*pi*ratio)
    elif shape == 3:
        return 1.0
    elif shape == 4:
        return 0.5 + 0.5*ratio
    elif shape == 5:
        if ratio <= 0.7:
            return 0.05 + 0.95 * ratio/0.7
        else:
            return 0.05 + 0.95 * (1.0 - ratio)/0.3
    elif shape == 6:
        return 1.0 - 0.8*ratio
    elif shape == 7:
        if ratio <= 0.7:
            return 0.5 + 0.5*ratio/0.7
        else:
            return 0.5 + 0.5*(1.0 - ratio)/0.3
    elif shape == 8:
        r = 1 - ratio
        if r == 1:
            v = custom[3]
        elif r >= custom[2]:
            pos = (r - custom[2]) / (1 - custom[2])
            #if (custom[0] >= custom[1] <= custom[3]) or (custom[0] <= custom[1] >= custom[3]):
            pos = pos * pos
            v = (pos * (custom[3] - custom[1])) + custom[1]
        else:
            pos = r / custom[2]
            #if (custom[0] >= custom[1] <= custom[3]) or (custom[0] <= custom[1] >= custom[3]):
            pos = 1 - (1 - pos) * (1 - pos)
            v = (pos * (custom[1] - custom[0])) + custom[0]

        return v

    elif shape == 9:
        if (ratio < (1 - pruneWidthPeak)) and (ratio > 0.0):
            return ((ratio/(1 - pruneWidthPeak))**prunePowerHigh)
        elif (ratio >= (1 - pruneWidthPeak)) and (ratio < 1.0):
            return (((1 - ratio)/pruneWidthPeak)**prunePowerLow)
        else:
            return 0.0

    elif shape == 10:
        return 0.5 + 0.5 * (1 - ratio)

# This function determines the actual number of splits at a given point using the global error
def splits(n):
    global splitError
    nEff = round(n + splitError, 0)
    splitError -= (nEff - n)
    return int(nEff)

def splits2(n):
    r = random()
    if r < n:
        return 1
    else:
        return 0

def splits3(n):
    ni = int(n)
    nf = n - int(n)
    r = random()
    if r < nf:
        return ni + 1
    else:
        return ni + 0

# Determine the declination from a given quaternion
def declination(quat):
    tempVec = zAxis.copy()
    tempVec.rotate(quat)
    tempVec.normalize()
    return degrees(acos(tempVec.z))

# Determines the angle of upward rotation of a segment due to attractUp
def curveUp(attractUp, quat, curveRes):
    tempVec = yAxis.copy()
    tempVec.rotate(quat)
    tempVec.normalize()

    dec = radians(declination(quat))
    curveUpAng = attractUp*dec*abs(tempVec.z)/curveRes
    if (-dec + curveUpAng) < -pi:
        curveUpAng = -pi + dec
    if (dec - curveUpAng) < 0:
        curveUpAng = dec
    return curveUpAng

# Evaluate a bezier curve for the parameter 0<=t<=1 along its length
def evalBez(p1, h1, h2, p2, t):
    return ((1-t)**3)*p1 + (3*t*(1-t)**2)*h1 + (3*(t**2)*(1-t))*h2 + (t**3)*p2

# Evaluate the unit tangent on a bezier curve for t
def evalBezTan(p1, h1, h2, p2, t):
    return ((-3*(1-t)**2)*p1 + (-6*t*(1-t) + 3*(1-t)**2)*h1 + (-3*(t**2) + 6*t*(1-t))*h2 + (3*t**2)*p2).normalized()

# Determine the range of t values along a splines length where child stems are formed
def findChildPoints(stemList, numChild):
    numPoints = sum([len(n.spline.bezier_points) for n in stemList])
    numSplines = len(stemList)
    numSegs = numPoints - numSplines
    numPerSeg = numChild/numSegs
    numMain = round(numPerSeg*stemList[0].segMax, 0)
    return [(a+1)/(numMain) for a in range(int(numMain))]

def findChildPoints2(stemList, numChild):
    return [(a+1)/(numChild) for a in range(int(numChild))]

# Find the coordinates, quaternion and radius for each t on the stem
def interpStem1(stem, tVals, lPar, parRad):
    points = stem.spline.bezier_points
    numPoints = len(points)
    checkVal = (stem.segMax - (numPoints - 1)) / stem.segMax
    # Loop through all the parametric values to be determined
    tempList = deque()
    for t in tVals:
        if t == 1.0:
            index = numPoints-2
            coord = points[-1].co
            quat = (points[-1].handle_right - points[-1].co).to_track_quat('Z', 'Y')
            radius = points[-1].radius

            tempList.append(childPoint(coord, quat, (parRad, radius), t, lPar, 'bone'+(str(stem.splN).rjust(3, '0'))+'.'+(str(index).rjust(3, '0'))))

        elif (t >= checkVal) and (t < 1.0):
            scaledT = (t - checkVal) / ((1 - checkVal) + .0001)
            length = (numPoints-1)*scaledT
            index = int(length)

            tTemp = length - index
            coord = evalBez(points[index].co, points[index].handle_right, points[index+1].handle_left, points[index+1].co, tTemp)
            quat = (evalBezTan(points[index].co, points[index].handle_right, points[index+1].handle_left, points[index+1].co, tTemp)).to_track_quat('Z', 'Y')
            radius = (1-tTemp)*points[index].radius + tTemp*points[index+1].radius # Not sure if this is the parent radius at the child point or parent start radius

            tempList.append(childPoint(coord, quat, (parRad, radius), t, lPar, 'bone'+(str(stem.splN).rjust(3, '0'))+'.'+(str(index).rjust(3, '0'))))

    return tempList

def interpStem(stem, tVals, lPar, parRad, maxOffset, baseSize):
    points = stem.spline.bezier_points
    numSegs = len(points) - 1
    stemLen = stem.segL * numSegs

    checkBottom = stem.offsetLen / maxOffset
    checkTop = checkBottom + (stemLen / maxOffset)

    # Loop through all the parametric values to be determined
    tempList = deque()
    for t in tVals:
        if (t >= checkBottom) and (t <= checkTop) and (t < 1.0):
            scaledT = (t - checkBottom) / (checkTop - checkBottom)
            ofst = ((t - baseSize) / (checkTop - baseSize)) * (1 - baseSize) + baseSize

            length = numSegs * scaledT
            index = int(length)
            tTemp = length - index

            coord = evalBez(points[index].co, points[index].handle_right, points[index+1].handle_left, points[index+1].co, tTemp)
            quat = (evalBezTan(points[index].co, points[index].handle_right, points[index+1].handle_left, points[index+1].co, tTemp)).to_track_quat('Z', 'Y')
            radius = (1-tTemp)*points[index].radius + tTemp*points[index+1].radius # Not sure if this is the parent radius at the child point or parent start radius

            tempList.append(childPoint(coord, quat, (parRad, radius), t, ofst, lPar, 'bone'+(str(stem.splN).rjust(3, '0'))+'.'+(str(index).rjust(3, '0'))))

    #add stem at tip
    index = numSegs-1
    coord = points[-1].co
    quat = (points[-1].handle_right - points[-1].co).to_track_quat('Z', 'Y')
    radius = points[-1].radius
    tempList.append(childPoint(coord, quat, (parRad, radius), 1, 1, lPar, 'bone'+(str(stem.splN).rjust(3, '0'))+'.'+(str(index).rjust(3, '0'))))

    return tempList

# round down bone number
def roundBone(bone, step):
    bone_i = bone[:-3]
    bone_n = int(bone[-3:])
    bone_n = int(bone_n / step) * step
    return bone_i + str(bone_n).rjust(3, '0')

# Convert a list of degrees to radians
def toRad(list):
    return [radians(a) for a in list]

def anglemean(a1, a2, fac):
    x1 = sin(a1)
    y1 = cos(a1)
    x2 = sin(a2)
    y2 = cos(a2)
    x = x1 + (x2 - x1) * fac
    y = y1 + (y2 - y1) * fac
    return atan2(x, y)


# This is the function which extends (or grows) a given stem.
def growSpline(n, stem, numSplit, splitAng, splitAngV, splineList, hType, splineToBone, closeTip, kp, splitHeight, outAtt, stemsegL,
               lenVar, taperCrown, boneStep, rotate, rotateV):

    #curv at base
    sCurv = stem.curv
    if (n == 0) and (kp <= splitHeight):
        sCurv = 0.0

    #curveangle = sCurv + (uniform(-stem.curvV, stem.curvV) * kp)
    #curveVar = uniform(-stem.curvV, stem.curvV) * kp
    curveangle = sCurv + (uniform(0, stem.curvV) * kp * stem.curvSignx)
    curveVar = uniform(0, stem.curvV) * kp * stem.curvSigny
    stem.curvSignx *= -1
    stem.curvSigny *= -1

    curveVarMat = Matrix.Rotation(curveVar, 3, 'Y')

    # First find the current direction of the stem
    dir = stem.quat()

    if n == 0:
        adir = zAxis.copy()
        adir.rotate(dir)

        ry = atan2(adir[0], adir[2])
        adir.rotate(Euler((0, -ry, 0)))
        rx = atan2(adir[1], adir[2])

        dir = Euler((-rx, ry, 0), 'XYZ')

    #length taperCrown
    if n == 0:
        dec = declination(dir) / 180
        dec = dec ** 2
        tf = 1 - (dec * taperCrown * 30)
        tf = max(.1, tf)
    else:
        tf = 1.0

    #outward attraction
    if (n > 0) and (kp > 0) and (outAtt > 0):
        p = stem.p.co.copy()
        d = atan2(p[0], -p[1]) + tau
        edir = dir.to_euler('XYZ', Euler((0, 0, d), 'XYZ'))
        d = anglemean(edir[2], d, (kp * outAtt))
        dirv = Euler((edir[0], edir[1], d), 'XYZ')
        dir = dirv.to_quaternion()

    #parent weight
#    parWeight = kp * degrees(stem.curvV) * pi
#    #parWeight = parWeight * kp
#    #parWeight = kp
#    if (n > 1) and (parWeight != 0):
#        d1 = zAxis.copy()
#        d2 = zAxis.copy()
#        d1.rotate(dir)
#        d2.rotate(stem.patentQuat)
#
#        x = d1[0] + ((d2[0] - d1[0]) * parWeight)
#        y = d1[1] + ((d2[1] - d1[1]) * parWeight)
#        z = d1[2] + ((d2[2] - d1[2]) * parWeight)
#
#        d3 = Vector((x, y, z))
#        dir = d3.to_track_quat('Z', 'Y')

    # If the stem splits, we need to add new splines etc
    if numSplit > 0:
        # Get the curve data
        cuData = stem.spline.id_data.name
        cu = bpy.data.curves[cuData]

        #calc split angles
        angle = choice([-1, 1]) * (splitAng + uniform(-splitAngV, splitAngV))
        if n > 0:
            #make branches flatter
            angle *= max(1 - declination(dir) / 90, 0) * .67 + .33
        spreadangle = choice([-1, 1]) * (splitAng + uniform(-splitAngV, splitAngV))

        #branchRotMat = Matrix.Rotation(radians(uniform(0, 360)), 3, 'Z')
        if not hasattr(stem, 'rLast'):
            stem.rLast = radians(uniform(0, 360))

        br = rotate[0] + uniform(-rotateV[0], rotateV[0])
        branchRot = stem.rLast + br
        branchRotMat = Matrix.Rotation(branchRot, 3, 'Z')
        stem.rLast = branchRot

        # Now for each split add the new spline and adjust the growth direction
        for i in range(numSplit):
            #find split scale
            lenV = uniform(1 - lenVar, 1 + lenVar)
            bScale = min(lenV * tf, 1)

            newSpline = cu.splines.new('BEZIER')
            newPoint = newSpline.bezier_points[-1]
            (newPoint.co, newPoint.handle_left_type, newPoint.handle_right_type) = (stem.p.co, 'VECTOR', 'VECTOR')
            newPoint.radius = (stem.radS*(1 - stem.seg/stem.segMax) + stem.radE*(stem.seg/stem.segMax)) * bScale
            # Here we make the new "sprouting" stems diverge from the current direction
            divRotMat = Matrix.Rotation(angle + curveangle, 3, 'X')
            dirVec = zAxis.copy()
            dirVec.rotate(divRotMat)

            #horizontal curvature variation
            dirVec.rotate(curveVarMat)

            if n == 0: #Special case for trunk splits
                dirVec.rotate(branchRotMat)

                ang = pi - ((tau) / (numSplit + 1)) * (i+1)
                dirVec.rotate(Matrix.Rotation(ang, 3, 'Z'))

            # Spread the stem out in a random fashion
            spreadMat = Matrix.Rotation(spreadangle, 3, 'Y')
            if n != 0: #Special case for trunk splits
                dirVec.rotate(spreadMat)

            dirVec.rotate(dir)

            # Introduce upward curvature
            upRotAxis = xAxis.copy()
            upRotAxis.rotate(dirVec.to_track_quat('Z', 'Y'))
            curveUpAng = curveUp(stem.vertAtt, dirVec.to_track_quat('Z', 'Y'), stem.segMax)
            upRotMat = Matrix.Rotation(-curveUpAng, 3, upRotAxis)
            dirVec.rotate(upRotMat)

            # Make the growth vec the length of a stem segment
            dirVec.normalize()

            #split length variation
            stemL = stemsegL * lenV
            dirVec *= stemL * tf
            ofst = stem.offsetLen + (stem.segL * (len(stem.spline.bezier_points) - 1))

            ##dirVec *= stem.segL

            # Get the end point position
            end_co = stem.p.co.copy()

            # Add the new point and adjust its coords, handles and radius
            newSpline.bezier_points.add()
            newPoint = newSpline.bezier_points[-1]
            (newPoint.co, newPoint.handle_left_type, newPoint.handle_right_type) = (end_co + dirVec, hType, hType)
            newPoint.radius = (stem.radS*(1 - (stem.seg + 1)/stem.segMax) + stem.radE*((stem.seg + 1)/stem.segMax)) * bScale
            if (stem.seg == stem.segMax-1) and closeTip:
                newPoint.radius = 0.0
            # If this isn't the last point on a stem, then we need to add it to the list of stems to continue growing
            #print(stem.seg != stem.segMax, stem.seg, stem.segMax)
            #if stem.seg != stem.segMax: # if probs not nessesary
            nstem = stemSpline(newSpline, stem.curv, stem.curvV, stem.vertAtt, stem.seg+1, stem.segMax, stemL, stem.children,
                               stem.radS * bScale, stem.radE * bScale, len(cu.splines)-1, ofst, stem.quat())
            nstem.splitlast = 1#numSplit #keep track of numSplit for next stem
            nstem.rLast = branchRot + pi
            splineList.append(nstem)
            bone = 'bone'+(str(stem.splN)).rjust(3, '0')+'.'+(str(len(stem.spline.bezier_points)-2)).rjust(3, '0')
            bone = roundBone(bone, boneStep[n])
            splineToBone.append((bone, False, True, len(stem.spline.bezier_points)-2))

        # The original spline also needs to keep growing so adjust its direction too
        divRotMat = Matrix.Rotation(-angle + curveangle, 3, 'X')
        dirVec = zAxis.copy()
        dirVec.rotate(divRotMat)

        #horizontal curvature variation
        dirVec.rotate(curveVarMat)

        if n == 0: #Special case for trunk splits
            dirVec.rotate(branchRotMat)

        #spread
        spreadMat = Matrix.Rotation(-spreadangle, 3, 'Y')
        if n != 0: #Special case for trunk splits
            dirVec.rotate(spreadMat)

        dirVec.rotate(dir)

        stem.splitlast = 1#numSplit #keep track of numSplit for next stem

    else:
        # If there are no splits then generate the growth direction without accounting for spreading of stems
        dirVec = zAxis.copy()
        divRotMat = Matrix.Rotation(curveangle, 3, 'X')
        dirVec.rotate(divRotMat)

        #horizontal curvature variation
        dirVec.rotate(curveVarMat)

        dirVec.rotate(dir)

        stem.splitlast = 0#numSplit #keep track of numSplit for next stem

    # Introduce upward curvature
    upRotAxis = xAxis.copy()
    upRotAxis.rotate(dirVec.to_track_quat('Z', 'Y'))
    curveUpAng = curveUp(stem.vertAtt, dirVec.to_track_quat('Z', 'Y'), stem.segMax)
    upRotMat = Matrix.Rotation(-curveUpAng, 3, upRotAxis)
    dirVec.rotate(upRotMat)

    dirVec.normalize()
    dirVec *= stem.segL * tf

    # Get the end point position
    end_co = stem.p.co.copy()

    stem.spline.bezier_points.add()
    newPoint = stem.spline.bezier_points[-1]
    (newPoint.co, newPoint.handle_left_type, newPoint.handle_right_type) = (end_co + dirVec, hType, hType)
    newPoint.radius = stem.radS*(1 - (stem.seg + 1)/stem.segMax) + stem.radE*((stem.seg + 1)/stem.segMax)
    if (stem.seg == stem.segMax-1) and closeTip:
        newPoint.radius = 0.0
    # There are some cases where a point cannot have handles as VECTOR straight away, set these now.
    if len(stem.spline.bezier_points) == 2:
        tempPoint = stem.spline.bezier_points[0]
        (tempPoint.handle_left_type, tempPoint.handle_right_type) = ('VECTOR', 'VECTOR')
    # Update the last point in the spline to be the newly added one
    stem.updateEnd()
    #return splineList

def genLeafMesh(leafScale, leafScaleX, leafScaleT, leafScaleV, loc, quat, offset, index, downAngle, downAngleV, rotate, rotateV, oldRot,
                bend, leaves, leafShape, leafangle, horzLeaves):
    if leafShape == 'hex':
        verts = [Vector((0, 0, 0)), Vector((0.5, 0, 1/3)), Vector((0.5, 0, 2/3)), Vector((0, 0, 1)), Vector((-0.5, 0, 2/3)), Vector((-0.5, 0, 1/3))]
        edges = [[0, 1], [1, 2], [2, 3], [3, 4], [4, 5], [5, 0], [0, 3]]
        faces = [[0, 1, 2, 3], [0, 3, 4, 5]]
    elif leafShape == 'rect':
        #verts = [Vector((1, 0, 0)), Vector((1, 0, 1)), Vector((-1, 0, 1)), Vector((-1, 0, 0))]
        verts = [Vector((.5, 0, 0)), Vector((.5, 0, 1)), Vector((-.5, 0, 1)), Vector((-.5, 0, 0))]
        edges = [[0, 1], [1, 2], [2, 3], [3, 0]]
        faces = [[0, 1, 2, 3]]
    elif leafShape == 'dFace':
        verts = [Vector((.5, .5, 0)), Vector((.5, -.5, 0)), Vector((-.5, -.5, 0)), Vector((-.5, .5, 0))]
        edges = [[0, 1], [1, 2], [2, 3], [3, 0]]
        faces = [[0, 3, 2, 1]]
    elif leafShape == 'dVert':
        verts = [Vector((0, 0, 1))]
        edges = []
        faces = []

    vertsList = []
    facesList = []
    normal = Vector((0, 0, 1))

    if leaves < 0:
        rotMat = Matrix.Rotation(oldRot, 3, 'Y')
    else:
        rotMat = Matrix.Rotation(oldRot, 3, 'Z')

    # If the -ve flag for rotate is used we need to find which side of the stem the last child point was and then grow in the opposite direction.
    if rotate < 0.0:
        oldRot = -copysign(rotate + uniform(-rotateV, rotateV), oldRot)
    else:
        # If the special -ve flag for leaves is used we need a different rotation of the leaf geometry
        if leaves == -1:
            #oldRot = 0
            rotMat = Matrix.Rotation(0, 3, 'Y')
        elif leaves < -1:
            oldRot += rotate / (-leaves - 1)
        else:
            oldRot += rotate + uniform(-rotateV, rotateV)

#    if leaves < 0:
#        rotMat = Matrix.Rotation(oldRot, 3, 'Y')
#    else:
#        rotMat = Matrix.Rotation(oldRot, 3, 'Z')

    if leaves >= 0:
        #downRotMat = Matrix.Rotation(downAngle+uniform(-downAngleV, downAngleV), 3, 'X')

        if downAngleV > 0.0:
            downV = -downAngleV * offset
        else:
            downV = uniform(-downAngleV, downAngleV)
        downRotMat = Matrix.Rotation(downAngle + downV, 3, 'X')

    #leaf scale variation
    if (leaves < -1) and (rotate != 0):
        f = 1 - abs((oldRot - (rotate / (-leaves - 1))) / (rotate / 2))
    else:
        f = offset

    if leafScaleT < 0:
        leafScale = leafScale * (1 - (1 - f) * -leafScaleT)
    else:
        leafScale = leafScale * (1 - f * leafScaleT)

    leafScale = leafScale * uniform(1 - leafScaleV, 1 + leafScaleV)

    if leafShape == 'dFace':
        leafScale = leafScale * .1

    # If the bending of the leaves is used we need to rotate them differently
    if (bend != 0.0) and (leaves >= 0):
        normal = yAxis.copy()
        orientationVec = zAxis.copy()

        normal.rotate(quat)
        orientationVec.rotate(quat)

        thetaPos = atan2(loc.y, loc.x)
        thetaBend = thetaPos - atan2(normal.y, normal.x)
        rotateZ = Matrix.Rotation(bend*thetaBend, 3, 'Z')
        normal.rotate(rotateZ)
        orientationVec.rotate(rotateZ)

        phiBend = atan2((normal.xy).length, normal.z)
        orientation = atan2(orientationVec.y, orientationVec.x)
        rotateZOrien = Matrix.Rotation(orientation, 3, 'X')

        rotateX = Matrix.Rotation(bend*phiBend, 3, 'Z')

        rotateZOrien2 = Matrix.Rotation(-orientation, 3, 'X')

    # For each of the verts we now rotate and scale them, then append them to the list to be added to the mesh
    for v in verts:
        v.z *= leafScale
        v.y *= leafScale
        v.x *= leafScaleX*leafScale

        v.rotate(Euler((0, 0, radians(180))))

        #leafangle
        v.rotate(Matrix.Rotation(radians(-leafangle), 3, 'X'))

        if rotate < 0:
            v.rotate(Euler((0, 0, radians(90))))
            if oldRot < 0:
                v.rotate(Euler((0, 0, radians(180))))

        if (leaves > 0) and (rotate > 0) and horzLeaves:
            nRotMat = Matrix.Rotation(-oldRot + rotate, 3, 'Z')
            v.rotate(nRotMat)

        if leaves > 0:
            v.rotate(downRotMat)

        v.rotate(rotMat)
        v.rotate(quat)

        if (bend != 0.0) and (leaves > 0):
            # Correct the rotation
            v.rotate(rotateZ)
            v.rotate(rotateZOrien)
            v.rotate(rotateX)
            v.rotate(rotateZOrien2)

    if leafShape == 'dVert':
        normal = verts[0]
        normal.normalize()
        v = loc
        vertsList.append([v.x, v.y, v.z])
    else:
        for v in verts:
            v += loc
            vertsList.append([v.x, v.y, v.z])
        for f in faces:
            facesList.append([f[0] + index, f[1] + index, f[2] + index, f[3] + index])

    return vertsList, facesList, normal, oldRot


def create_armature(armAnim, leafP, cu, frameRate, leafMesh, leafObj, leafVertSize, leaves, levelCount, splineToBone,
                    treeOb, wind, gust, gustF, af1, af2, af3, leafAnim, loopFrames, previewArm, armLevels, makeMesh, boneStep):
    arm = bpy.data.armatures.new('tree')
    armOb = bpy.data.objects.new('treeArm', arm)
    bpy.context.scene.objects.link(armOb)
    # Create a new action to store all animation
    newAction = bpy.data.actions.new(name='windAction')
    armOb.animation_data_create()
    armOb.animation_data.action = newAction
    arm.draw_type = 'STICK'
    arm.use_deform_delay = True
    # Add the armature modifier to the curve
    armMod = treeOb.modifiers.new('windSway', 'ARMATURE')
    if previewArm:
        armMod.show_viewport = False
        arm.draw_type = 'WIRE'
        treeOb.hide = True
    armMod.use_apply_on_spline = True
    armMod.object = armOb
    armMod.use_bone_envelopes = True
    armMod.use_vertex_groups = False # curves don't have vertex groups (yet)
    # If there are leaves then they need a modifier
    if leaves:
        armMod = leafObj.modifiers.new('windSway', 'ARMATURE')
        armMod.object = armOb
        armMod.use_bone_envelopes = False
        armMod.use_vertex_groups = True
    # Make sure all objects are deselected (may not be required?)
    for ob in bpy.data.objects:
        ob.select = False

    fps = bpy.context.scene.render.fps
    animSpeed = (24 / fps) * frameRate

    # Set the armature as active and go to edit mode to add bones
    bpy.context.scene.objects.active = armOb
    bpy.ops.object.mode_set(mode='EDIT')
    # For all the splines in the curve we need to add bones at each bezier point
    for i, parBone in enumerate(splineToBone):
        if (i < levelCount[armLevels]) or (armLevels == -1) or (not makeMesh):
            s = cu.splines[i]
            b = None
            # Get some data about the spline like length and number of points
            numPoints = len(s.bezier_points) - 1

            #find branching level
            level = 0
            for l, c in enumerate(levelCount):
                if i < c:
                    level = l
                    break
            level = min(level, 3)

            step = boneStep[level]

            # Calculate things for animation
            if armAnim:
                splineL = numPoints * ((s.bezier_points[0].co - s.bezier_points[1].co).length)
                # Set the random phase difference of the animation
                bxOffset = uniform(0, tau)
                byOffset = uniform(0, tau)
                # Set the phase multiplier for the spline
                #bMult_r = (s.bezier_points[0].radius / max(splineL, 1e-6)) * (1 / 15) * (1 / frameRate)
                #bMult = degrees(bMult_r)  # This shouldn't have to be in degrees but it looks much better in animation
                bMult = (1 / max(splineL ** .5, 1e-6)) * (1 / 4)
                #print((1 / bMult) * tau) #print wavelength in frames

                windFreq1 = bMult * animSpeed
                windFreq2 = 0.7 * bMult * animSpeed
                if loopFrames != 0:
                    bMult_l = 1 / (loopFrames / tau)
                    fRatio = max(1, round(windFreq1 / bMult_l))
                    fgRatio = max(1, round(windFreq2 / bMult_l))
                    windFreq1 = fRatio * bMult_l
                    windFreq2 = fgRatio * bMult_l

            # For all the points in the curve (less the last) add a bone and name it by the spline it will affect
            nx = 0
            for n in range(0, numPoints, step):
                oldBone = b
                boneName = 'bone' + (str(i)).rjust(3, '0') + '.' + (str(n)).rjust(3, '0')
                b = arm.edit_bones.new(boneName)
                b.head = s.bezier_points[n].co
                nx += step
                nx = min(nx, numPoints)
                b.tail = s.bezier_points[nx].co

                b.head_radius = s.bezier_points[n].radius
                b.tail_radius = s.bezier_points[n + 1].radius
                b.envelope_distance = 0.001

#                # If there are leaves then we need a new vertex group so they will attach to the bone
#                if not leafAnim:
#                    if (len(levelCount) > 1) and (i >= levelCount[-2]) and leafObj:
#                        leafObj.vertex_groups.new(boneName)
#                    elif (len(levelCount) == 1) and leafObj:
#                        leafObj.vertex_groups.new(boneName)

                # If this is first point of the spline then it must be parented to the level above it
                if n == 0:
                    if parBone:
                        b.parent = arm.edit_bones[parBone]
                # Otherwise, we need to attach it to the previous bone in the spline
                else:
                    b.parent = oldBone
                    b.use_connect = True
                # If there isn't a previous bone then it shouldn't be attached
                if not oldBone:
                    b.use_connect = False

                # Add the animation to the armature if required
                if armAnim:
                    # Define all the required parameters of the wind sway by the dimension of the spline
                    #a0 = 4 * splineL * (1 - n / (numPoints + 1)) / max(s.bezier_points[n].radius, 1e-6)
                    a0 = 2 * (splineL / numPoints) * (1 - n / (numPoints + 1)) / max(s.bezier_points[n].radius, 1e-6)
                    a0 = a0 * min(step, numPoints)
                    #a0 = (splineL / numPoints) / max(s.bezier_points[n].radius, 1e-6)
                    a1 = (wind / 50) * a0
                    a2 = a1 * .65  #(windGust / 50) * a0 + a1 / 2

                    p = s.bezier_points[nx].co - s.bezier_points[n].co
                    p.normalize()
                    ag = (wind * gust / 50) * a0
                    a3 = -p[0] * ag
                    a4 = p[2] * ag

                    a1 = radians(a1)
                    a2 = radians(a2)
                    a3 = radians(a3)
                    a4 = radians(a4)

                    #wind bending
                    if loopFrames == 0:
                        swayFreq = gustF * (tau / fps) * frameRate  #animSpeed # .075 # 0.02
                    else:
                        swayFreq = 1 / (loopFrames / tau)

                    # Prevent tree base from rotating
                    if (boneName == "bone000.000") or (boneName == "bone000.001"):
                        a1 = 0
                        a2 = 0
                        a3 = 0
                        a4 = 0

                    # Add new fcurves for each sway as well as the modifiers
                    swayX = armOb.animation_data.action.fcurves.new('pose.bones["' + boneName + '"].rotation_euler', 0)
                    swayY = armOb.animation_data.action.fcurves.new('pose.bones["' + boneName + '"].rotation_euler', 2)

                    swayXMod1 = swayX.modifiers.new(type='FNGENERATOR')
                    swayXMod2 = swayX.modifiers.new(type='FNGENERATOR')

                    swayYMod1 = swayY.modifiers.new(type='FNGENERATOR')
                    swayYMod2 = swayY.modifiers.new(type='FNGENERATOR')

                    # Set the parameters for each modifier
                    swayXMod1.amplitude = a1
                    swayXMod1.phase_offset = bxOffset
                    swayXMod1.phase_multiplier = windFreq1

                    swayXMod2.amplitude = a2
                    swayXMod2.phase_offset = 0.7 * bxOffset
                    swayXMod2.phase_multiplier = windFreq2
                    swayXMod2.use_additive = True

                    swayYMod1.amplitude = a1
                    swayYMod1.phase_offset = byOffset
                    swayYMod1.phase_multiplier = windFreq1

                    swayYMod2.amplitude = a2
                    swayYMod2.phase_offset = 0.7 * byOffset
                    swayYMod2.phase_multiplier = windFreq2
                    swayYMod2.use_additive = True

                    #wind bending
                    swayYMod3 = swayY.modifiers.new(type='FNGENERATOR')
                    swayYMod3.amplitude = a3
                    swayYMod3.phase_multiplier = swayFreq
                    swayYMod3.value_offset = .6 * a3
                    swayYMod3.use_additive = True

                    swayXMod3 = swayX.modifiers.new(type='FNGENERATOR')
                    swayXMod3.amplitude = a4
                    swayXMod3.phase_multiplier = swayFreq
                    swayXMod3.value_offset = .6 * a4
                    swayXMod3.use_additive = True

    if leaves:
        bonelist = [b.name for b in arm.edit_bones]
        vertexGroups = OrderedDict()
        for i, cp in enumerate(leafP):
            # find leafs parent bone
            leafParent = roundBone(cp.parBone, boneStep[armLevels])
            idx = int(leafParent[4:-4])
            while leafParent not in bonelist:
                #find parent bone of parent bone
                leafParent = splineToBone[idx]
                idx = int(leafParent[4:-4])

            if leafAnim:
                bname = "leaf" + str(i)
                b = arm.edit_bones.new(bname)
                b.head = cp.co
                b.tail = cp.co + Vector((0, 0, .02))
                b.envelope_distance = 0.0
                b.parent = arm.edit_bones[leafParent]

                vertexGroups[bname] = [v.index for v in leafMesh.vertices[leafVertSize * i:(leafVertSize * i + leafVertSize)]]

                if armAnim:
                    # Define all the required parameters of the wind sway by the dimension of the spline
                    a1 = wind * .25
                    a1 *= af1

                    bMult = (1 / animSpeed) * 6
                    bMult *= 1 / max(af2, .001)

                    ofstRand = af3
                    bxOffset = uniform(-ofstRand, ofstRand)
                    byOffset = uniform(-ofstRand, ofstRand)

                    # Add new fcurves for each sway as well as the modifiers
                    swayX = armOb.animation_data.action.fcurves.new('pose.bones["' + bname + '"].rotation_euler', 0)
                    swayY = armOb.animation_data.action.fcurves.new('pose.bones["' + bname + '"].rotation_euler', 2)

                    # Add keyframe so noise works
                    swayX.keyframe_points.add()
                    swayY.keyframe_points.add()
                    swayX.keyframe_points[0].co = (0, 0)
                    swayY.keyframe_points[0].co = (0, 0)

                    # Add noise modifiers
                    swayXMod = swayX.modifiers.new(type='NOISE')
                    swayYMod = swayY.modifiers.new(type='NOISE')

                    if loopFrames != 0:
                        swayXMod.use_restricted_range = True
                        swayXMod.frame_end = loopFrames
                        swayXMod.blend_in = 4
                        swayXMod.blend_out = 4
                        swayYMod.use_restricted_range = True
                        swayYMod.frame_end = loopFrames
                        swayYMod.blend_in = 4
                        swayYMod.blend_out = 4

                    swayXMod.scale = bMult
                    swayXMod.strength = a1
                    swayXMod.offset = bxOffset

                    swayYMod.scale = bMult
                    swayYMod.strength = a1
                    swayYMod.offset = byOffset

            else:
                if leafParent not in vertexGroups:
                    vertexGroups[leafParent] = []
                vertexGroups[leafParent].extend([v.index for v in leafMesh.vertices[leafVertSize * i:(leafVertSize * i + leafVertSize)]])

        for group in vertexGroups:
            leafObj.vertex_groups.new(group)
            leafObj.vertex_groups[group].add(vertexGroups[group], 1.0, 'ADD')

    # Now we need the rotation mode to be 'XYZ' to ensure correct rotation
    bpy.ops.object.mode_set(mode='OBJECT')
    for p in armOb.pose.bones:
        p.rotation_mode = 'XYZ'
    treeOb.parent = armOb


def kickstart_trunk(addstem, levels, leaves, branches, cu, curve, curveRes, curveV, attractUp, length, lengthV, ratio, ratioPower, resU, scale0, scaleV0,
                    scaleVal, taper, minRadius, rootFlare):
    newSpline = cu.splines.new('BEZIER')
    cu.resolution_u = resU
    newPoint = newSpline.bezier_points[-1]
    newPoint.co = Vector((0, 0, 0))
    newPoint.handle_right = Vector((0, 0, 1))
    newPoint.handle_left = Vector((0, 0, -1))
    # (newPoint.handle_right_type, newPoint.handle_left_type) = ('VECTOR', 'VECTOR')
    branchL = scaleVal * length[0]
    curveVal = curve[0] / curveRes[0]
    #curveVal = curveVal * (branchL / scaleVal)
    if levels == 1:
        childStems = leaves
    else:
        childStems = branches[1]
    startRad = scaleVal * ratio * scale0 * uniform(1-scaleV0, 1+scaleV0) ## * (scale0 + uniform(-scaleV0, scaleV0)) #
    endRad = (startRad * (1 - taper[0])) ** ratioPower
    startRad = max(startRad, minRadius)
    endRad = max(endRad, minRadius)
    newPoint.radius = startRad * rootFlare
    addstem(
        stemSpline(newSpline, curveVal, curveV[0] / curveRes[0], attractUp[0], 0, curveRes[0], branchL / curveRes[0],
                   childStems, startRad, endRad, 0, 0, None))


def fabricate_stems(addsplinetobone, addstem, baseSize, branches, childP, cu, curve, curveBack, curveRes, curveV, attractUp,
                    downAngle, downAngleV, leafDist, leaves, length, lengthV, levels, n, ratioPower, resU,
                    rotate, rotateV, scaleVal, shape, storeN, taper, shapeS, minRadius, radiusTweak, customShape, rMode, segSplits,
                    useOldDownAngle, useParentAngle, boneStep):

    #prevent baseSize from going to 1.0
    baseSize = min(0.999, baseSize)

    # Store the old rotation to allow new stems to be rotated away from the previous one.
    oldRotate = 0

    #use fancy child point selection / rotation
    if (n == 1) and (rMode != "original"):
        childP_T = OrderedDict()
        childP_L = []
        for p in childP:
            if p.offset == 1:
                childP_L.append(p)
            else:
                if p.offset not in childP_T:
                    childP_T[p.offset] = [p]
                else:
                    childP_T[p.offset].append(p)

        childP_T = [childP_T[k] for k in sorted(childP_T.keys())]

        childP = []
        rot_a = []
        for p in childP_T:
            if rMode == "rotate":
                if rotate[n] < 0.0:
                    oldRotate = -copysign(rotate[n], oldRotate)
                else:
                    oldRotate += rotate[n]
                bRotate = oldRotate + uniform(-rotateV[n], rotateV[n])

                #choose start point whose angle is closest to the rotate angle
                a1 = bRotate % tau
                a_diff = []
                for a in p:
                    a2 = atan2(a.co[0], -a.co[1])
                    d = min((a1-a2+tau)%tau, (a2-a1+tau)%tau)
                    a_diff.append(d)

                idx = a_diff.index(min(a_diff))

                #find actual rotate angle from branch location
                br = p[idx]
                b = br.co
                vx = sin(bRotate)
                vy = cos(bRotate)
                v = Vector((vx, vy))

                bD = ((b[0] * b[0] + b[1] * b[1]) ** .5)
                bL = br.lengthPar * length[1] * shapeRatio(shape, (1 - br.offset) / (1 - baseSize), custom=customShape)

                #account for down angle
                if downAngleV[1] > 0: