add_curve_sapling/utils.py

   1 # -*- coding: utf-8 -*-
   2 # ##### BEGIN GPL LICENSE BLOCK #####
   3 #
   4 #  This program is free software; you can redistribute it and/or
   5 #  modify it under the terms of the GNU General Public License
   6 #  as published by the Free Software Foundation; either version 2
   7 #  of the License, or (at your option) any later version.
   8 #
   9 #  This program is distributed in the hope that it will be useful,
  10 #  but WITHOUT ANY WARRANTY; without even the implied warranty of
  11 #  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  12 #  GNU General Public License for more details.
  13 #
  14 #  You should have received a copy of the GNU General Public License
  15 #  along with this program; if not, write to the Free Software Foundation,
  16 #  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
  17 #
  18 # ##### END GPL LICENSE BLOCK #####
  19
  20
  21 import bpy
  22
  23 import time
  24 import copy
  25
  26 from mathutils import (
  27         Euler,
  28         Matrix,
  29         Vector,
  30         )
  31 from math import pi, sin, degrees, radians, atan2, copysign, cos, acos
  32 from math import floor
  33 from random import random, uniform, seed, choice, getstate, setstate, randint
  34 from collections import deque, OrderedDict
  35
  36 tau = 2 * pi
  37
  38 # Initialise the split error and axis vectors
  39 splitError = 0.0
  40 zAxis = Vector((0, 0, 1))
  41 yAxis = Vector((0, 1, 0))
  42 xAxis = Vector((1, 0, 0))
  43
  44
  45 # This class will contain a part of the tree which needs to be extended and the required tree parameters
  46 class stemSpline:
  47     def __init__(self, spline, curvature, curvatureV, attractUp, segments, maxSegs,
  48                  segLength, childStems, stemRadStart, stemRadEnd, splineNum, ofst, pquat):
  49         self.spline = spline
  50         self.p = spline.bezier_points[-1]
  51         self.curv = curvature
  52         self.curvV = curvatureV
  53         self.vertAtt = attractUp
  54         self.seg = segments
  55         self.segMax = maxSegs
  56         self.segL = segLength
  57         self.children = childStems
  58         self.radS = stemRadStart
  59         self.radE = stemRadEnd
  60         self.splN = splineNum
  61         self.offsetLen = ofst
  62         self.patentQuat = pquat
  63         self.curvSignx = 1
  64         self.curvSigny = 1
  65
  66     # This method determines the quaternion of the end of the spline
  67     def quat(self):
  68         if len(self.spline.bezier_points) == 1:
  69             return ((self.spline.bezier_points[-1].handle_right -
  70                      self.spline.bezier_points[-1].co).normalized()).to_track_quat('Z', 'Y')
  71         else:
  72             return ((self.spline.bezier_points[-1].co -
  73                      self.spline.bezier_points[-2].co).normalized()).to_track_quat('Z', 'Y')
  74
  75     # Determine the declination
  76     def dec(self):
  77         tempVec = zAxis.copy()
  78         tempVec.rotate(self.quat())
  79         return zAxis.angle(tempVec)
  80
  81     # Update the end of the spline and increment the segment count
  82     def updateEnd(self):
  83         self.p = self.spline.bezier_points[-1]
  84         self.seg += 1
  85
  86
  87 # This class contains the data for a point where a new branch will sprout
  88 class childPoint:
  89     def __init__(self, coords, quat, radiusPar, offset, sOfst, lengthPar, parBone):
  90         self.co = coords
  91         self.quat = quat
  92         self.radiusPar = radiusPar
  93         self.offset = offset
  94         self.stemOffset = sOfst
  95         self.lengthPar = lengthPar
  96         self.parBone = parBone
  97
  98
  99 # This function calculates the shape ratio as defined in the paper
 100 def shapeRatio(shape, ratio, pruneWidthPeak=0.0, prunePowerHigh=0.0, prunePowerLow=0.0, custom=None):
 101     if shape == 0:
 102         return 0.05 + 0.95 * ratio  # 0.2 + 0.8 * ratio
 103     elif shape == 1:
 104         return 0.2 + 0.8 * sin(pi * ratio)
 105     elif shape == 2:
 106         return 0.2 + 0.8 * sin(0.5 * pi * ratio)
 107     elif shape == 3:
 108         return 1.0
 109     elif shape == 4:
 110         return 0.5 + 0.5 * ratio
 111     elif shape == 5:
 112         if ratio <= 0.7:
 113             return 0.05 + 0.95 * ratio / 0.7
 114         else:
 115             return 0.05 + 0.95 * (1.0 - ratio) / 0.3
 116     elif shape == 6:
 117         return 1.0 - 0.8 * ratio
 118     elif shape == 7:
 119         if ratio <= 0.7:
 120             return 0.5 + 0.5 * ratio / 0.7
 121         else:
 122             return 0.5 + 0.5 * (1.0 - ratio) / 0.3
 123     elif shape == 8:
 124         r = 1 - ratio
 125         if r == 1:
 126             v = custom[3]
 127         elif r >= custom[2]:
 128             pos = (r - custom[2]) / (1 - custom[2])
 129             # if (custom[0] >= custom[1] <= custom[3]) or (custom[0] <= custom[1] >= custom[3]):
 130             pos = pos * pos
 131             v = (pos * (custom[3] - custom[1])) + custom[1]
 132         else:
 133             pos = r / custom[2]
 134             # if (custom[0] >= custom[1] <= custom[3]) or (custom[0] <= custom[1] >= custom[3]):
 135             pos = 1 - (1 - pos) * (1 - pos)
 136             v = (pos * (custom[1] - custom[0])) + custom[0]
 137
 138         return v
 139
 140     elif shape == 9:
 141         if (ratio < (1 - pruneWidthPeak)) and (ratio > 0.0):
 142             return ((ratio / (1 - pruneWidthPeak))**prunePowerHigh)
 143         elif (ratio >= (1 - pruneWidthPeak)) and (ratio < 1.0):
 144             return (((1 - ratio) / pruneWidthPeak)**prunePowerLow)
 145         else:
 146             return 0.0
 147
 148     elif shape == 10:
 149         return 0.5 + 0.5 * (1 - ratio)
 150
 151
 152 # This function determines the actual number of splits at a given point using the global error
 153 def splits(n):
 154     global splitError
 155     nEff = round(n + splitError, 0)
 156     splitError -= (nEff - n)
 157     return int(nEff)
 158
 159
 160 def splits2(n):
 161     r = random()
 162     if r < n:
 163         return 1
 164     else:
 165         return 0
 166
 167
 168 def splits3(n):
 169     ni = int(n)
 170     nf = n - int(n)
 171     r = random()
 172     if r < nf:
 173         return ni + 1
 174     else:
 175         return ni + 0
 176
 177
 178 # Determine the declination from a given quaternion
 179 def declination(quat):
 180     tempVec = zAxis.copy()
 181     tempVec.rotate(quat)
 182     tempVec.normalize()
 183     return degrees(acos(tempVec.z))
 184
 185
 186 # Determines the angle of upward rotation of a segment due to attractUp
 187 def curveUp(attractUp, quat, curveRes):
 188     tempVec = yAxis.copy()
 189     tempVec.rotate(quat)
 190     tempVec.normalize()
 191
 192     dec = radians(declination(quat))
 193     curveUpAng = attractUp * dec * abs(tempVec.z) / curveRes
 194     if (-dec + curveUpAng) < -pi:
 195         curveUpAng = -pi + dec
 196     if (dec - curveUpAng) < 0:
 197         curveUpAng = dec
 198     return curveUpAng
 199
 200
 201 # Evaluate a bezier curve for the parameter 0<=t<=1 along its length
 202 def evalBez(p1, h1, h2, p2, t):
 203     return ((1 - t)**3) * p1 + (3 * t * (1 - t)**2) * h1 + (3 * (t**2) * (1 - t)) * h2 + (t**3) * p2
 204
 205
 206 # Evaluate the unit tangent on a bezier curve for t
 207 def evalBezTan(p1, h1, h2, p2, t):
 208     return (
 209             (-3 * (1 - t)**2) * p1 + (-6 * t * (1 - t) + 3 * (1 - t)**2) * h1 +
 210             (-3 * (t**2) + 6 * t * (1 - t)) * h2 + (3 * t**2) * p2
 211             ).normalized()
 212
 213
 214 # Determine the range of t values along a splines length where child stems are formed
 215 def findChildPoints(stemList, numChild):
 216     numPoints = sum([len(n.spline.bezier_points) for n in stemList])
 217     numSplines = len(stemList)
 218     numSegs = numPoints - numSplines
 219     numPerSeg = numChild / numSegs
 220     numMain = round(numPerSeg * stemList[0].segMax, 0)
 221     return [(a + 1) / (numMain) for a in range(int(numMain))]
 222
 223
 224 def findChildPoints2(stemList, numChild):
 225     return [(a + 1) / (numChild) for a in range(int(numChild))]
 226
 227
 228 # Find the coordinates, quaternion and radius for each t on the stem
 229 def interpStem1(stem, tVals, lPar, parRad):
 230     points = stem.spline.bezier_points
 231     numPoints = len(points)
 232     checkVal = (stem.segMax - (numPoints - 1)) / stem.segMax
 233     # Loop through all the parametric values to be determined
 234     tempList = deque()
 235     for t in tVals:
 236         if t == 1.0:
 237             index = numPoints - 2
 238             coord = points[-1].co
 239             quat = (points[-1].handle_right - points[-1].co).to_track_quat('Z', 'Y')
 240             radius = points[-1].radius
 241
 242             tempList.append(
 243                     childPoint(coord, quat, (parRad, radius), t, lPar, 'bone' +
 244                               (str(stem.splN).rjust(3, '0')) + '.' + (str(index).rjust(3, '0')))
 245                     )
 246
 247         elif (t >= checkVal) and (t < 1.0):
 248             scaledT = (t - checkVal) / ((1 - checkVal) + .0001)
 249             length = (numPoints - 1) * scaledT
 250             index = int(length)
 251
 252             tTemp = length - index
 253             coord = evalBez(
 254                         points[index].co, points[index].handle_right,
 255                         points[index + 1].handle_left, points[index + 1].co, tTemp
 256                         )
 257             quat = (
 258                 evalBezTan(
 259                     points[index].co, points[index].handle_right,
 260                     points[index + 1].handle_left, points[index + 1].co, tTemp)
 261                     ).to_track_quat('Z', 'Y')
 262             # Not sure if this is the parent radius at the child point or parent start radius
 263             radius = (1 - tTemp) * points[index].radius + tTemp * points[index + 1].radius
 264
 265             tempList.append(
 266                     childPoint(
 267                             coord, quat, (parRad, radius), t, lPar, 'bone' +
 268                             (str(stem.splN).rjust(3, '0')) + '.' + (str(index).rjust(3, '0')))
 269                             )
 270     return tempList
 271
 272
 273 def interpStem(stem, tVals, lPar, parRad, maxOffset, baseSize):
 274     points = stem.spline.bezier_points
 275     numSegs = len(points) - 1
 276     stemLen = stem.segL * numSegs
 277
 278     checkBottom = stem.offsetLen / maxOffset
 279     checkTop = checkBottom + (stemLen / maxOffset)
 280
 281     # Loop through all the parametric values to be determined
 282     tempList = deque()
 283     for t in tVals:
 284         if (t >= checkBottom) and (t <= checkTop) and (t < 1.0):
 285             scaledT = (t - checkBottom) / (checkTop - checkBottom)
 286             ofst = ((t - baseSize) / (checkTop - baseSize)) * (1 - baseSize) + baseSize
 287
 288             length = numSegs * scaledT
 289             index = int(length)
 290             tTemp = length - index
 291
 292             coord = evalBez(
 293                         points[index].co, points[index].handle_right,
 294                         points[index + 1].handle_left, points[index + 1].co, tTemp
 295                         )
 296             quat = (
 297                 evalBezTan(
 298                     points[index].co, points[index].handle_right,
 299                     points[index + 1].handle_left, points[index + 1].co, tTemp
 300                     )
 301                 ).to_track_quat('Z', 'Y')
 302             # Not sure if this is the parent radius at the child point or parent start radius
 303             radius = (1 - tTemp) * points[index].radius + tTemp * points[index + 1].radius
 304
 305             tempList.append(
 306                     childPoint(
 307                         coord, quat, (parRad, radius), t, ofst, lPar,
 308                         'bone' + (str(stem.splN).rjust(3, '0')) + '.' + (str(index).rjust(3, '0')))
 309                     )
 310
 311     # add stem at tip
 312     index = numSegs - 1
 313     coord = points[-1].co
 314     quat = (points[-1].handle_right - points[-1].co).to_track_quat('Z', 'Y')
 315     radius = points[-1].radius
 316     tempList.append(
 317                 childPoint(
 318                         coord, quat, (parRad, radius), 1, 1, lPar,
 319                         'bone' + (str(stem.splN).rjust(3, '0')) + '.' + (str(index).rjust(3, '0'))
 320                         )
 321                     )
 322
 323     return tempList
 324
 325
 326 # round down bone number
 327 def roundBone(bone, step):
 328     bone_i = bone[:-3]
 329     bone_n = int(bone[-3:])
 330     bone_n = int(bone_n / step) * step
 331     return bone_i + str(bone_n).rjust(3, '0')
 332
 333
 334 # Convert a list of degrees to radians
 335 def toRad(list):
 336     return [radians(a) for a in list]
 337
 338
 339 def anglemean(a1, a2, fac):
 340     x1 = sin(a1)
 341     y1 = cos(a1)
 342     x2 = sin(a2)
 343     y2 = cos(a2)
 344     x = x1 + (x2 - x1) * fac
 345     y = y1 + (y2 - y1) * fac
 346     return atan2(x, y)
 347
 348
 349 # This is the function which extends (or grows) a given stem.
 350 def growSpline(n, stem, numSplit, splitAng, splitAngV, splineList,
 351                hType, splineToBone, closeTip, kp, splitHeight, outAtt, stemsegL,
 352                lenVar, taperCrown, boneStep, rotate, rotateV):
 353
 354     # curv at base
 355     sCurv = stem.curv
 356     if (n == 0) and (kp <= splitHeight):
 357         sCurv = 0.0
 358
 359     # curveangle = sCurv + (uniform(-stem.curvV, stem.curvV) * kp)
 360     # curveVar = uniform(-stem.curvV, stem.curvV) * kp
 361     curveangle = sCurv + (uniform(0, stem.curvV) * kp * stem.curvSignx)
 362     curveVar = uniform(0, stem.curvV) * kp * stem.curvSigny
 363     stem.curvSignx *= -1
 364     stem.curvSigny *= -1
 365
 366     curveVarMat = Matrix.Rotation(curveVar, 3, 'Y')
 367
 368     # First find the current direction of the stem
 369     dir = stem.quat()
 370
 371     if n == 0:
 372         adir = zAxis.copy()
 373         adir.rotate(dir)
 374
 375         ry = atan2(adir[0], adir[2])
 376         adir.rotate(Euler((0, -ry, 0)))
 377         rx = atan2(adir[1], adir[2])
 378
 379         dir = Euler((-rx, ry, 0), 'XYZ')
 380
 381     # length taperCrown
 382     if n == 0:
 383         dec = declination(dir) / 180
 384         dec = dec ** 2
 385         tf = 1 - (dec * taperCrown * 30)
 386         tf = max(.1, tf)
 387     else:
 388         tf = 1.0
 389
 390     # outward attraction
 391     if (n > 0) and (kp > 0) and (outAtt > 0):
 392         p = stem.p.co.copy()
 393         d = atan2(p[0], -p[1]) + tau
 394         edir = dir.to_euler('XYZ', Euler((0, 0, d), 'XYZ'))
 395         d = anglemean(edir[2], d, (kp * outAtt))
 396         dirv = Euler((edir[0], edir[1], d), 'XYZ')
 397         dir = dirv.to_quaternion()
 398     """
 399     # parent weight
 400     parWeight = kp * degrees(stem.curvV) * pi
 401     parWeight = parWeight * kp
 402     parWeight = kp
 403     if (n > 1) and (parWeight != 0):
 404         d1 = zAxis.copy()
 405         d2 = zAxis.copy()
 406         d1.rotate(dir)
 407         d2.rotate(stem.patentQuat)
 408
 409         x = d1[0] + ((d2[0] - d1[0]) * parWeight)
 410         y = d1[1] + ((d2[1] - d1[1]) * parWeight)
 411         z = d1[2] + ((d2[2] - d1[2]) * parWeight)
 412
 413         d3 = Vector((x, y, z))
 414         dir = d3.to_track_quat('Z', 'Y')
 415     """
 416
 417     # If the stem splits, we need to add new splines etc
 418     if numSplit > 0:
 419         # Get the curve data
 420         cuData = stem.spline.id_data.name
 421         cu = bpy.data.curves[cuData]
 422
 423         # calc split angles
 424         angle = choice([-1, 1]) * (splitAng + uniform(-splitAngV, splitAngV))
 425         if n > 0:
 426             # make branches flatter
 427             angle *= max(1 - declination(dir) / 90, 0) * .67 + .33
 428         spreadangle = choice([-1, 1]) * (splitAng + uniform(-splitAngV, splitAngV))
 429
 430         # branchRotMat = Matrix.Rotation(radians(uniform(0, 360)), 3, 'Z')
 431         if not hasattr(stem, 'rLast'):
 432             stem.rLast = radians(uniform(0, 360))
 433
 434         br = rotate[0] + uniform(-rotateV[0], rotateV[0])
 435         branchRot = stem.rLast + br
 436         branchRotMat = Matrix.Rotation(branchRot, 3, 'Z')
 437         stem.rLast = branchRot
 438
 439         # Now for each split add the new spline and adjust the growth direction
 440         for i in range(numSplit):
 441             # find split scale
 442             lenV = uniform(1 - lenVar, 1 + lenVar)
 443             bScale = min(lenV * tf, 1)
 444
 445             newSpline = cu.splines.new('BEZIER')
 446             newPoint = newSpline.bezier_points[-1]
 447             (newPoint.co, newPoint.handle_left_type, newPoint.handle_right_type) = (stem.p.co, 'VECTOR', 'VECTOR')
 448             newPoint.radius = (
 449                         stem.radS * (1 - stem.seg / stem.segMax) + stem.radE * (stem.seg / stem.segMax)
 450                         ) * bScale
 451             # Here we make the new "sprouting" stems diverge from the current direction
 452             divRotMat = Matrix.Rotation(angle + curveangle, 3, 'X')
 453             dirVec = zAxis.copy()
 454             dirVec.rotate(divRotMat)
 455
 456             # horizontal curvature variation
 457             dirVec.rotate(curveVarMat)
 458
 459             if n == 0:  # Special case for trunk splits
 460                 dirVec.rotate(branchRotMat)
 461
 462                 ang = pi - ((tau) / (numSplit + 1)) * (i + 1)
 463                 dirVec.rotate(Matrix.Rotation(ang, 3, 'Z'))
 464
 465             # Spread the stem out in a random fashion
 466             spreadMat = Matrix.Rotation(spreadangle, 3, 'Y')
 467             if n != 0:  # Special case for trunk splits
 468                 dirVec.rotate(spreadMat)
 469
 470             dirVec.rotate(dir)
 471
 472             # Introduce upward curvature
 473             upRotAxis = xAxis.copy()
 474             upRotAxis.rotate(dirVec.to_track_quat('Z', 'Y'))
 475             curveUpAng = curveUp(stem.vertAtt, dirVec.to_track_quat('Z', 'Y'), stem.segMax)
 476             upRotMat = Matrix.Rotation(-curveUpAng, 3, upRotAxis)
 477             dirVec.rotate(upRotMat)
 478
 479             # Make the growth vec the length of a stem segment
 480             dirVec.normalize()
 481
 482             # split length variation
 483             stemL = stemsegL * lenV
 484             dirVec *= stemL * tf
 485             ofst = stem.offsetLen + (stem.segL * (len(stem.spline.bezier_points) - 1))
 486
 487             # dirVec *= stem.segL
 488
 489             # Get the end point position
 490             end_co = stem.p.co.copy()
 491
 492             # Add the new point and adjust its coords, handles and radius
 493             newSpline.bezier_points.add(1)
 494             newPoint = newSpline.bezier_points[-1]
 495             (newPoint.co, newPoint.handle_left_type, newPoint.handle_right_type) = (end_co + dirVec, hType, hType)
 496             newPoint.radius = (
 497                         stem.radS * (1 - (stem.seg + 1) / stem.segMax) +
 498                         stem.radE * ((stem.seg + 1) / stem.segMax)
 499                         ) * bScale
 500             if (stem.seg == stem.segMax - 1) and closeTip:
 501                 newPoint.radius = 0.0
 502             # If this isn't the last point on a stem, then we need to add it
 503             # to the list of stems to continue growing
 504             # print(stem.seg != stem.segMax, stem.seg, stem.segMax)
 505             # if stem.seg != stem.segMax: # if probs not necessary
 506             nstem = stemSpline(
 507                         newSpline, stem.curv, stem.curvV, stem.vertAtt, stem.seg + 1,
 508                         stem.segMax, stemL, stem.children,
 509                         stem.radS * bScale, stem.radE * bScale, len(cu.splines) - 1, ofst, stem.quat()
 510                         )
 511             nstem.splitlast = 1  # numSplit  # keep track of numSplit for next stem
 512             nstem.rLast = branchRot + pi
 513             splineList.append(nstem)
 514             bone = 'bone' + (str(stem.splN)).rjust(3, '0') + '.' + \
 515                     (str(len(stem.spline.bezier_points) - 2)).rjust(3, '0')
 516             bone = roundBone(bone, boneStep[n])
 517             splineToBone.append((bone, False, True, len(stem.spline.bezier_points) - 2))
 518
 519         # The original spline also needs to keep growing so adjust its direction too
 520         divRotMat = Matrix.Rotation(-angle + curveangle, 3, 'X')
 521         dirVec = zAxis.copy()
 522         dirVec.rotate(divRotMat)
 523
 524         # horizontal curvature variation
 525         dirVec.rotate(curveVarMat)
 526
 527         if n == 0:  # Special case for trunk splits
 528             dirVec.rotate(branchRotMat)
 529
 530         # spread
 531         spreadMat = Matrix.Rotation(-spreadangle, 3, 'Y')
 532         if n != 0:  # Special case for trunk splits
 533             dirVec.rotate(spreadMat)
 534
 535         dirVec.rotate(dir)
 536
 537         stem.splitlast = 1  # numSplit #keep track of numSplit for next stem
 538
 539     else:
 540         # If there are no splits then generate the growth direction without accounting for spreading of stems
 541         dirVec = zAxis.copy()
 542         divRotMat = Matrix.Rotation(curveangle, 3, 'X')
 543         dirVec.rotate(divRotMat)
 544
 545         # horizontal curvature variation
 546         dirVec.rotate(curveVarMat)
 547
 548         dirVec.rotate(dir)
 549
 550         stem.splitlast = 0  # numSplit #keep track of numSplit for next stem
 551
 552     # Introduce upward curvature
 553     upRotAxis = xAxis.copy()
 554     upRotAxis.rotate(dirVec.to_track_quat('Z', 'Y'))
 555     curveUpAng = curveUp(stem.vertAtt, dirVec.to_track_quat('Z', 'Y'), stem.segMax)
 556     upRotMat = Matrix.Rotation(-curveUpAng, 3, upRotAxis)
 557     dirVec.rotate(upRotMat)
 558
 559     dirVec.normalize()
 560     dirVec *= stem.segL * tf
 561
 562     # Get the end point position
 563     end_co = stem.p.co.copy()
 564
 565     stem.spline.bezier_points.add(1)
 566     newPoint = stem.spline.bezier_points[-1]
 567     (newPoint.co, newPoint.handle_left_type, newPoint.handle_right_type) = (end_co + dirVec, hType, hType)
 568     newPoint.radius = stem.radS * (1 - (stem.seg + 1) / stem.segMax) + \
 569                       stem.radE * ((stem.seg + 1) / stem.segMax)
 570
 571     if (stem.seg == stem.segMax - 1) and closeTip:
 572         newPoint.radius = 0.0
 573     # There are some cases where a point cannot have handles as VECTOR straight away, set these now
 574     if len(stem.spline.bezier_points) == 2:
 575         tempPoint = stem.spline.bezier_points[0]
 576         (tempPoint.handle_left_type, tempPoint.handle_right_type) = ('VECTOR', 'VECTOR')
 577     # Update the last point in the spline to be the newly added one
 578     stem.updateEnd()
 579     # return splineList
 580
 581
 582 def genLeafMesh(leafScale, leafScaleX, leafScaleT, leafScaleV, loc, quat,
 583                 offset, index, downAngle, downAngleV, rotate, rotateV, oldRot,
 584                 bend, leaves, leafShape, leafangle, horzLeaves):
 585     if leafShape == 'hex':
 586         verts = [
 587             Vector((0, 0, 0)), Vector((0.5, 0, 1 / 3)), Vector((0.5, 0, 2 / 3)),
 588             Vector((0, 0, 1)), Vector((-0.5, 0, 2 / 3)), Vector((-0.5, 0, 1 / 3))
 589             ]
 590         edges = [[0, 1], [1, 2], [2, 3], [3, 4], [4, 5], [5, 0], [0, 3]]
 591         faces = [[0, 1, 2, 3], [0, 3, 4, 5]]
 592     elif leafShape == 'rect':
 593         # verts = [Vector((1, 0, 0)), Vector((1, 0, 1)), Vector((-1, 0, 1)), Vector((-1, 0, 0))]
 594         verts = [Vector((.5, 0, 0)), Vector((.5, 0, 1)), Vector((-.5, 0, 1)), Vector((-.5, 0, 0))]
 595         edges = [[0, 1], [1, 2], [2, 3], [3, 0]]
 596         faces = [[0, 1, 2, 3]]
 597     elif leafShape == 'dFace':
 598         verts = [Vector((.5, .5, 0)), Vector((.5, -.5, 0)), Vector((-.5, -.5, 0)), Vector((-.5, .5, 0))]
 599         edges = [[0, 1], [1, 2], [2, 3], [3, 0]]
 600         faces = [[0, 3, 2, 1]]
 601     elif leafShape == 'dVert':
 602         verts = [Vector((0, 0, 1))]
 603         edges = []
 604         faces = []
 605
 606     vertsList = []
 607     facesList = []
 608     normal = Vector((0, 0, 1))
 609
 610     if leaves < 0:
 611         rotMat = Matrix.Rotation(oldRot, 3, 'Y')
 612     else:
 613         rotMat = Matrix.Rotation(oldRot, 3, 'Z')
 614
 615     # If the -ve flag for rotate is used we need to find which side of the stem
 616     # the last child point was and then grow in the opposite direction
 617     if rotate < 0.0:
 618         oldRot = -copysign(rotate + uniform(-rotateV, rotateV), oldRot)
 619     else:
 620         # If the special -ve flag for leaves is used we need a different rotation of the leaf geometry
 621         if leaves == -1:
 622             # oldRot = 0
 623             rotMat = Matrix.Rotation(0, 3, 'Y')
 624         elif leaves < -1:
 625             oldRot += rotate / (-leaves - 1)
 626         else:
 627             oldRot += rotate + uniform(-rotateV, rotateV)
 628     """
 629     if leaves < 0:
 630         rotMat = Matrix.Rotation(oldRot, 3, 'Y')
 631     else:
 632         rotMat = Matrix.Rotation(oldRot, 3, 'Z')
 633     """
 634     if leaves >= 0:
 635         # downRotMat = Matrix.Rotation(downAngle+uniform(-downAngleV, downAngleV), 3, 'X')
 636
 637         if downAngleV > 0.0:
 638             downV = -downAngleV * offset
 639         else:
 640             downV = uniform(-downAngleV, downAngleV)
 641         downRotMat = Matrix.Rotation(downAngle + downV, 3, 'X')
 642
 643     # leaf scale variation
 644     if (leaves < -1) and (rotate != 0):
 645         f = 1 - abs((oldRot - (rotate / (-leaves - 1))) / (rotate / 2))
 646     else:
 647         f = offset
 648
 649     if leafScaleT < 0:
 650         leafScale = leafScale * (1 - (1 - f) * -leafScaleT)
 651     else:
 652         leafScale = leafScale * (1 - f * leafScaleT)
 653
 654     leafScale = leafScale * uniform(1 - leafScaleV, 1 + leafScaleV)
 655
 656     if leafShape == 'dFace':
 657         leafScale = leafScale * .1
 658
 659     # If the bending of the leaves is used we need to rotate them differently
 660     if (bend != 0.0) and (leaves >= 0):
 661         normal = yAxis.copy()
 662         orientationVec = zAxis.copy()
 663
 664         normal.rotate(quat)
 665         orientationVec.rotate(quat)
 666
 667         thetaPos = atan2(loc.y, loc.x)
 668         thetaBend = thetaPos - atan2(normal.y, normal.x)
 669         rotateZ = Matrix.Rotation(bend * thetaBend, 3, 'Z')
 670         normal.rotate(rotateZ)
 671         orientationVec.rotate(rotateZ)
 672
 673         phiBend = atan2((normal.xy).length, normal.z)
 674         orientation = atan2(orientationVec.y, orientationVec.x)
 675         rotateZOrien = Matrix.Rotation(orientation, 3, 'X')
 676
 677         rotateX = Matrix.Rotation(bend * phiBend, 3, 'Z')
 678
 679         rotateZOrien2 = Matrix.Rotation(-orientation, 3, 'X')
 680
 681     # For each of the verts we now rotate and scale them, then append them to the list to be added to the mesh
 682     for v in verts:
 683         v.z *= leafScale
 684         v.y *= leafScale
 685         v.x *= leafScaleX * leafScale
 686
 687         v.rotate(Euler((0, 0, radians(180))))
 688
 689         # leafangle
 690         v.rotate(Matrix.Rotation(radians(-leafangle), 3, 'X'))
 691
 692         if rotate < 0:
 693             v.rotate(Euler((0, 0, radians(90))))
 694             if oldRot < 0:
 695                 v.rotate(Euler((0, 0, radians(180))))
 696
 697         if (leaves > 0) and (rotate > 0) and horzLeaves:
 698             nRotMat = Matrix.Rotation(-oldRot + rotate, 3, 'Z')
 699             v.rotate(nRotMat)
 700
 701         if leaves > 0:
 702             v.rotate(downRotMat)
 703
 704         v.rotate(rotMat)
 705         v.rotate(quat)
 706
 707         if (bend != 0.0) and (leaves > 0):
 708             # Correct the rotation
 709             v.rotate(rotateZ)
 710             v.rotate(rotateZOrien)
 711             v.rotate(rotateX)
 712             v.rotate(rotateZOrien2)
 713
 714     if leafShape == 'dVert':
 715         normal = verts[0]
 716         normal.normalize()
 717         v = loc
 718         vertsList.append([v.x, v.y, v.z])
 719     else:
 720         for v in verts:
 721             v += loc
 722             vertsList.append([v.x, v.y, v.z])
 723         for f in faces:
 724             facesList.append([f[0] + index, f[1] + index, f[2] + index, f[3] + index])
 725
 726     return vertsList, facesList, normal, oldRot
 727
 728
 729 def create_armature(armAnim, leafP, cu, frameRate, leafMesh, leafObj, leafVertSize, leaves,
 730                     levelCount, splineToBone, treeOb, wind, gust, gustF, af1, af2, af3,
 731                     leafAnim, loopFrames, previewArm, armLevels, makeMesh, boneStep):
 732     arm = bpy.data.armatures.new('tree')
 733     armOb = bpy.data.objects.new('treeArm', arm)
 734     bpy.context.scene.collection.objects.link(armOb)
 735     # Create a new action to store all animation
 736     newAction = bpy.data.actions.new(name='windAction')
 737     armOb.animation_data_create()
 738     armOb.animation_data.action = newAction
 739     arm.display_type = 'STICK'
 740     arm.use_deform_delay = True
 741     # Add the armature modifier to the curve
 742     armMod = treeOb.modifiers.new('windSway', 'ARMATURE')
 743     if previewArm:
 744         armMod.show_viewport = False
 745         arm.display_type = 'WIRE'
 746         treeOb.hide_viewport = True
 747     armMod.use_apply_on_spline = True
 748     armMod.object = armOb
 749     armMod.use_bone_envelopes = True
 750     armMod.use_vertex_groups = False  # curves don't have vertex groups (yet)
 751     # If there are leaves then they need a modifier
 752     if leaves:
 753         armMod = leafObj.modifiers.new('windSway', 'ARMATURE')
 754         armMod.object = armOb
 755         armMod.use_bone_envelopes = False
 756         armMod.use_vertex_groups = True
 757
 758     # Make sure all objects are deselected (may not be required?)
 759     for ob in bpy.data.objects:
 760         ob.select_set(state=False)
 761
 762     fps = bpy.context.scene.render.fps
 763     animSpeed = (24 / fps) * frameRate
 764
 765     # Set the armature as active and go to edit mode to add bones
 766     bpy.context.view_layer.objects.active = armOb
 767     bpy.ops.object.mode_set(mode='EDIT')
 768     # For all the splines in the curve we need to add bones at each bezier point
 769     for i, parBone in enumerate(splineToBone):
 770         if (i < levelCount[armLevels]) or (armLevels == -1) or (not makeMesh):
 771             s = cu.splines[i]
 772             b = None
 773             # Get some data about the spline like length and number of points
 774             numPoints = len(s.bezier_points) - 1
 775
 776             # find branching level
 777             level = 0
 778             for l, c in enumerate(levelCount):
 779                 if i < c:
 780                     level = l
 781                     break
 782             level = min(level, 3)
 783
 784             step = boneStep[level]
 785
 786             # Calculate things for animation
 787             if armAnim:
 788                 splineL = numPoints * ((s.bezier_points[0].co - s.bezier_points[1].co).length)
 789                 # Set the random phase difference of the animation
 790                 bxOffset = uniform(0, tau)
 791                 byOffset = uniform(0, tau)
 792                 # Set the phase multiplier for the spline
 793                 # bMult_r = (s.bezier_points[0].radius / max(splineL, 1e-6)) * (1 / 15) * (1 / frameRate)
 794                 # This shouldn't have to be in degrees but it looks much better in animation
 795                 # bMult = degrees(bMult_r)
 796                 bMult = (1 / max(splineL ** .5, 1e-6)) * (1 / 4)
 797                 # print((1 / bMult) * tau) #print wavelength in frames
 798
 799                 windFreq1 = bMult * animSpeed
 800                 windFreq2 = 0.7 * bMult * animSpeed
 801                 if loopFrames != 0:
 802                     bMult_l = 1 / (loopFrames / tau)
 803                     fRatio = max(1, round(windFreq1 / bMult_l))
 804                     fgRatio = max(1, round(windFreq2 / bMult_l))
 805                     windFreq1 = fRatio * bMult_l
 806                     windFreq2 = fgRatio * bMult_l
 807
 808             # For all the points in the curve (less the last) add a bone and name it by the spline it will affect
 809             nx = 0
 810             for n in range(0, numPoints, step):
 811                 oldBone = b
 812                 boneName = 'bone' + (str(i)).rjust(3, '0') + '.' + (str(n)).rjust(3, '0')
 813                 b = arm.edit_bones.new(boneName)
 814                 b.head = s.bezier_points[n].co
 815                 nx += step
 816                 nx = min(nx, numPoints)
 817                 b.tail = s.bezier_points[nx].co
 818
 819                 b.head_radius = s.bezier_points[n].radius
 820                 b.tail_radius = s.bezier_points[n + 1].radius
 821                 b.envelope_distance = 0.001
 822                 """
 823                 # If there are leaves then we need a new vertex group so they will attach to the bone
 824                 if not leafAnim:
 825                     if (len(levelCount) > 1) and (i >= levelCount[-2]) and leafObj:
 826                         leafObj.vertex_groups.new(name=boneName)
 827                     elif (len(levelCount) == 1) and leafObj:
 828                         leafObj.vertex_groups.new(name=boneName)
 829                 """
 830                 # If this is first point of the spline then it must be parented to the level above it
 831                 if n == 0:
 832                     if parBone:
 833                         b.parent = arm.edit_bones[parBone]
 834                 # Otherwise, we need to attach it to the previous bone in the spline
 835                 else:
 836                     b.parent = oldBone
 837                     b.use_connect = True
 838                 # If there isn't a previous bone then it shouldn't be attached
 839                 if not oldBone:
 840                     b.use_connect = False
 841
 842                 # Add the animation to the armature if required
 843                 if armAnim:
 844                     # Define all the required parameters of the wind sway by the dimension of the spline
 845                     # a0 = 4 * splineL * (1 - n / (numPoints + 1)) / max(s.bezier_points[n].radius, 1e-6)
 846                     a0 = 2 * (splineL / numPoints) * (1 - n / (numPoints + 1)) / max(s.bezier_points[n].radius, 1e-6)
 847                     a0 = a0 * min(step, numPoints)
 848                     # a0 = (splineL / numPoints) / max(s.bezier_points[n].radius, 1e-6)
 849                     a1 = (wind / 50) * a0
 850                     a2 = a1 * .65  # (windGust / 50) * a0 + a1 / 2
 851
 852                     p = s.bezier_points[nx].co - s.bezier_points[n].co
 853                     p.normalize()
 854                     ag = (wind * gust / 50) * a0
 855                     a3 = -p[0] * ag
 856                     a4 = p[2] * ag
 857
 858                     a1 = radians(a1)
 859                     a2 = radians(a2)
 860                     a3 = radians(a3)
 861                     a4 = radians(a4)
 862
 863                     # wind bending
 864                     if loopFrames == 0:
 865                         swayFreq = gustF * (tau / fps) * frameRate  # animSpeed # .075 # 0.02
 866                     else:
 867                         swayFreq = 1 / (loopFrames / tau)
 868
 869                     # Prevent tree base from rotating
 870                     if (boneName == "bone000.000") or (boneName == "bone000.001"):
 871                         a1 = 0
 872                         a2 = 0
 873                         a3 = 0
 874                         a4 = 0
 875
 876                     # Add new fcurves for each sway as well as the modifiers
 877                     swayX = armOb.animation_data.action.fcurves.new(
 878                                             'pose.bones["' + boneName + '"].rotation_euler', index=0
 879                                             )
 880                     swayY = armOb.animation_data.action.fcurves.new(
 881                                             'pose.bones["' + boneName + '"].rotation_euler', index=2
 882                                             )
 883                     swayXMod1 = swayX.modifiers.new(type='FNGENERATOR')
 884                     swayXMod2 = swayX.modifiers.new(type='FNGENERATOR')
 885
 886                     swayYMod1 = swayY.modifiers.new(type='FNGENERATOR')
 887                     swayYMod2 = swayY.modifiers.new(type='FNGENERATOR')
 888
 889                     # Set the parameters for each modifier
 890                     swayXMod1.amplitude = a1
 891                     swayXMod1.phase_offset = bxOffset
 892                     swayXMod1.phase_multiplier = windFreq1
 893
 894                     swayXMod2.amplitude = a2
 895                     swayXMod2.phase_offset = 0.7 * bxOffset
 896                     swayXMod2.phase_multiplier = windFreq2
 897                     swayXMod2.use_additive = True
 898
 899                     swayYMod1.amplitude = a1
 900                     swayYMod1.phase_offset = byOffset
 901                     swayYMod1.phase_multiplier = windFreq1
 902
 903                     swayYMod2.amplitude = a2
 904                     swayYMod2.phase_offset = 0.7 * byOffset
 905                     swayYMod2.phase_multiplier = windFreq2
 906                     swayYMod2.use_additive = True
 907
 908                     # wind bending
 909                     swayYMod3 = swayY.modifiers.new(type='FNGENERATOR')
 910                     swayYMod3.amplitude = a3
 911                     swayYMod3.phase_multiplier = swayFreq
 912                     swayYMod3.value_offset = .6 * a3
 913                     swayYMod3.use_additive = True
 914
 915                     swayXMod3 = swayX.modifiers.new(type='FNGENERATOR')
 916                     swayXMod3.amplitude = a4
 917                     swayXMod3.phase_multiplier = swayFreq
 918                     swayXMod3.value_offset = .6 * a4
 919                     swayXMod3.use_additive = True
 920
 921     if leaves:
 922         bonelist = [b.name for b in arm.edit_bones]
 923         vertexGroups = OrderedDict()
 924         for i, cp in enumerate(leafP):
 925             # find leafs parent bone
 926             leafParent = roundBone(cp.parBone, boneStep[armLevels])
 927             idx = int(leafParent[4:-4])
 928             while leafParent not in bonelist:
 929                 # find parent bone of parent bone
 930                 leafParent = splineToBone[idx]
 931                 idx = int(leafParent[4:-4])
 932
 933             if leafAnim:
 934                 bname = "leaf" + str(i)
 935                 b = arm.edit_bones.new(bname)
 936                 b.head = cp.co
 937                 b.tail = cp.co + Vector((0, 0, .02))
 938                 b.envelope_distance = 0.0
 939                 b.parent = arm.edit_bones[leafParent]
 940
 941                 vertexGroups[bname] = [
 942                                     v.index for v in
 943                                     leafMesh.vertices[leafVertSize * i:(leafVertSize * i + leafVertSize)]
 944                                     ]
 945
 946                 if armAnim:
 947                     # Define all the required parameters of the wind sway by the dimension of the spline
 948                     a1 = wind * .25
 949                     a1 *= af1
 950
 951                     bMult = (1 / animSpeed) * 6
 952                     bMult *= 1 / max(af2, .001)
 953
 954                     ofstRand = af3
 955                     bxOffset = uniform(-ofstRand, ofstRand)
 956                     byOffset = uniform(-ofstRand, ofstRand)
 957
 958                     # Add new fcurves for each sway as well as the modifiers
 959                     swayX = armOb.animation_data.action.fcurves.new(
 960                                                 'pose.bones["' + bname + '"].rotation_euler', index=0
 961                                                 )
 962                     swayY = armOb.animation_data.action.fcurves.new(
 963                                                 'pose.bones["' + bname + '"].rotation_euler', index=2
 964                                                 )
 965                     # Add keyframe so noise works
 966                     swayX.keyframe_points.add(1)
 967                     swayY.keyframe_points.add(1)
 968                     swayX.keyframe_points[0].co = (0, 0)
 969                     swayY.keyframe_points[0].co = (0, 0)
 970
 971                     # Add noise modifiers
 972                     swayXMod = swayX.modifiers.new(type='NOISE')
 973                     swayYMod = swayY.modifiers.new(type='NOISE')
 974
 975                     if loopFrames != 0:
 976                         swayXMod.use_restricted_range = True
 977                         swayXMod.frame_end = loopFrames
 978                         swayXMod.blend_in = 4
 979                         swayXMod.blend_out = 4
 980                         swayYMod.use_restricted_range = True
 981                         swayYMod.frame_end = loopFrames
 982                         swayYMod.blend_in = 4
 983                         swayYMod.blend_out = 4
 984
 985                     swayXMod.scale = bMult
 986                     swayXMod.strength = a1
 987                     swayXMod.offset = bxOffset
 988
 989                     swayYMod.scale = bMult
 990                     swayYMod.strength = a1
 991                     swayYMod.offset = byOffset
 992
 993             else:
 994                 if leafParent not in vertexGroups:
 995                     vertexGroups[leafParent] = []
 996                 vertexGroups[leafParent].extend(
 997                                         [v.index for v in
 998                                         leafMesh.vertices[leafVertSize * i:(leafVertSize * i + leafVertSize)]]
 999                                         )
1000
1001         for group in vertexGroups:
1002             leafObj.vertex_groups.new(name=group)
1003             leafObj.vertex_groups[group].add(vertexGroups[group], 1.0, 'ADD')
1004
1005     # Now we need the rotation mode to be 'XYZ' to ensure correct rotation
1006     bpy.ops.object.mode_set(mode='OBJECT')
1007     for p in armOb.pose.bones:
1008         p.rotation_mode = 'XYZ'
1009     treeOb.parent = armOb
1010
1011
1012 def kickstart_trunk(addstem, levels, leaves, branches, cu, curve, curveRes,
1013                     curveV, attractUp, length, lengthV, ratio, ratioPower,
1014                     resU, scale0, scaleV0, scaleVal, taper, minRadius, rootFlare):
1015     newSpline = cu.splines.new('BEZIER')
1016     cu.resolution_u = resU
1017     newPoint = newSpline.bezier_points[-1]
1018     newPoint.co = Vector((0, 0, 0))
1019     newPoint.handle_right = Vector((0, 0, 1))
1020     newPoint.handle_left = Vector((0, 0, -1))
1021     # (newPoint.handle_right_type, newPoint.handle_left_type) = ('VECTOR', 'VECTOR')
1022     branchL = scaleVal * length[0]
1023     curveVal = curve[0] / curveRes[0]
1024     # curveVal = curveVal * (branchL / scaleVal)
1025     if levels == 1:
1026         childStems = leaves
1027     else:
1028         childStems = branches[1]
1029     startRad = scaleVal * ratio * scale0 * uniform(1 - scaleV0, 1 + scaleV0)  # * (scale0 + uniform(-scaleV0, scaleV0))
1030     endRad = (startRad * (1 - taper[0])) ** ratioPower
1031     startRad = max(startRad, minRadius)
1032     endRad = max(endRad, minRadius)
1033     newPoint.radius = startRad * rootFlare
1034     addstem(
1035         stemSpline(
1036                 newSpline, curveVal, curveV[0] / curveRes[0], attractUp[0],
1037                 0, curveRes[0], branchL / curveRes[0],
1038                 childStems, startRad, endRad, 0, 0, None
1039                 )
1040             )
1041
1042
1043 def fabricate_stems(addsplinetobone, addstem, baseSize, branches, childP, cu, curve, curveBack,
1044                     curveRes, curveV, attractUp, downAngle, downAngleV, leafDist, leaves, length,
1045                     lengthV, levels, n, ratioPower, resU, rotate, rotateV, scaleVal, shape, storeN,
1046                     taper, shapeS, minRadius, radiusTweak, customShape, rMode, segSplits,
1047                     useOldDownAngle, useParentAngle, boneStep):
1048
1049     # prevent baseSize from going to 1.0
1050     baseSize = min(0.999, baseSize)
1051
1052     # Store the old rotation to allow new stems to be rotated away from the previous one.
1053     oldRotate = 0
1054
1055     # use fancy child point selection / rotation
1056     if (n == 1) and (rMode != "original"):
1057         childP_T = OrderedDict()
1058         childP_L = []
1059         for p in childP:
1060             if p.offset == 1:
1061                 childP_L.append(p)
1062             else:
1063                 if p.offset not in childP_T:
1064                     childP_T[p.offset] = [p]
1065                 else:
1066                     childP_T[p.offset].append(p)
1067
1068         childP_T = [childP_T[k] for k in sorted(childP_T.keys())]
1069
1070         childP = []
1071         rot_a = []
1072         for p in childP_T:
1073             if rMode == "rotate":
1074                 if rotate[n] < 0.0:
1075                     oldRotate = -copysign(rotate[n], oldRotate)
1076                 else:
1077                     oldRotate += rotate[n]
1078                 bRotate = oldRotate + uniform(-rotateV[n], rotateV[n])
1079
1080                 # choose start point whose angle is closest to the rotate angle
1081                 a1 = bRotate % tau
1082                 a_diff = []
1083                 for a in p:
1084                     a2 = atan2(a.co[0], -a.co[1])
1085                     d = min((a1 - a2 + tau) % tau, (a2 - a1 + tau) % tau)
1086                     a_diff.append(d)
1087
1088                 idx = a_diff.index(min(a_diff))
1089
1090                 # find actual rotate angle from branch location
1091                 br = p[idx]
1092                 b = br.co
1093                 vx = sin(bRotate)
1094                 vy = cos(bRotate)
1095                 v = Vector((vx, vy))
1096
1097                 bD = ((b[0] * b[0] + b[1] * b[1]) ** .5)
1098                 bL = br.lengthPar * length[1] * shapeRatio(shape, (1 - br.offset) / (1 - baseSize), custom=customShape)
1099
1100                 # account for down angle
1101                 if downAngleV[1] > 0:
1102                     downA = downAngle[n] + (-downAngleV[n] * (1 - (1 - br.offset) / (1 - baseSize)) ** 2)
1103                 else:
1104                     downA = downAngle[n]
1105                 if downA < (.5 * pi):
1106                     downA = sin(downA) ** 2
1107                     bL *= downA
1108
1109                 bL *= 0.33
1110                 v *= (bD + bL)
1111
1112                 bv = Vector((b[0], -b[1]))
1113                 cv = v - bv
1114                 a = atan2(cv[0], cv[1])
1115                 # rot_a.append(a)
1116                 """
1117                 # add fill points at top  #experimental
1118                 fillHeight = 1 - degrees(rotateV[3]) # 0.8
1119                 if fillHeight < 1:
1120                     w = (p[0].offset - fillHeight) / (1- fillHeight)
1121                     prob_b = random() < w
1122                 else:
1123                     prob_b = False
1124
1125                 if (p[0].offset > fillHeight): # prob_b and (len(p) > 1):  ##(p[0].offset > fillHeight) and
1126                     childP.append(p[randint(0, len(p)-1)])
1127                     rot_a.append(bRotate)# + pi)
1128                 """
1129                 childP.append(p[idx])
1130                 rot_a.append(a)
1131
1132             else:
1133                 idx = randint(0, len(p) - 1)
1134                 childP.append(p[idx])
1135             # childP.append(p[idx])
1136
1137         childP.extend(childP_L)
1138         rot_a.extend([0] * len(childP_L))
1139
1140         oldRotate = 0
1141
1142     for i, p in enumerate(childP):
1143         # Add a spline and set the coordinate of the first point.
1144         newSpline = cu.splines.new('BEZIER')
1145         cu.resolution_u = resU
1146         newPoint = newSpline.bezier_points[-1]
1147         newPoint.co = p.co
1148         tempPos = zAxis.copy()
1149         # If the -ve flag for downAngle is used we need a special formula to find it
1150         if useOldDownAngle:
1151             if downAngleV[n] < 0.0:
1152                 downV = downAngleV[n] * (1 - 2 * (.2 + .8 * ((1 - p.offset) / (1 - baseSize))))
1153             # Otherwise just find a random value
1154             else:
1155                 downV = uniform(-downAngleV[n], downAngleV[n])
1156         else:
1157             if downAngleV[n] < 0.0:
1158                 downV = uniform(-downAngleV[n], downAngleV[n])
1159             else:
1160                 downV = -downAngleV[n] * (1 - (1 - p.offset) / (1 - baseSize)) ** 2  # (110, 80) = (60, -50)
1161
1162         if p.offset == 1:
1163             downRotMat = Matrix.Rotation(0, 3, 'X')
1164         else:
1165             downRotMat = Matrix.Rotation(downAngle[n] + downV, 3, 'X')
1166
1167         # If the -ve flag for rotate is used we need to find which side of the stem
1168         # the last child point was and then grow in the opposite direction
1169         if rotate[n] < 0.0:
1170             oldRotate = -copysign(rotate[n], oldRotate)
1171         # Otherwise just generate a random number in the specified range
1172         else:
1173             oldRotate += rotate[n]
1174         bRotate = oldRotate + uniform(-rotateV[n], rotateV[n])
1175
1176         if (n == 1) and (rMode == "rotate"):
1177             bRotate = rot_a[i]
1178
1179         rotMat = Matrix.Rotation(bRotate, 3, 'Z')
1180
1181         # Rotate the direction of growth and set the new point coordinates
1182         tempPos.rotate(downRotMat)
1183         tempPos.rotate(rotMat)
1184
1185         # use quat angle
1186         if (rMode == "rotate") and (n == 1) and (p.offset != 1):
1187             if useParentAngle:
1188                 edir = p.quat.to_euler('XYZ', Euler((0, 0, bRotate), 'XYZ'))
1189                 edir[0] = 0
1190                 edir[1] = 0
1191
1192                 edir[2] = -edir[2]
1193                 tempPos.rotate(edir)
1194
1195                 dec = declination(p.quat)
1196                 tempPos.rotate(Matrix.Rotation(radians(dec), 3, 'X'))
1197
1198                 edir[2] = -edir[2]
1199                 tempPos.rotate(edir)
1200         else:
1201             tempPos.rotate(p.quat)
1202
1203         newPoint.handle_right = p.co + tempPos
1204
1205         # Make length variation inversely proportional to segSplits
1206         # lenV = (1 - min(segSplits[n], 1)) * lengthV[n]
1207
1208         # Find branch length and the number of child stems.
1209         maxbL = scaleVal
1210         for l in length[:n + 1]:
1211             maxbL *= l
1212         lMax = length[n]  # * uniform(1 - lenV, 1 + lenV)
1213         if n == 1:
1214             lShape = shapeRatio(shape, (1 - p.stemOffset) / (1 - baseSize), custom=customShape)
1215         else:
1216             lShape = shapeRatio(shapeS, (1 - p.stemOffset) / (1 - baseSize))
1217         branchL = p.lengthPar * lMax * lShape
1218         childStems = branches[min(3, n + 1)] * (0.1 + 0.9 * (branchL / maxbL))
1219
1220         # If this is the last level before leaves then we need to generate the child points differently
1221         if (storeN == levels - 1):
1222             if leaves < 0:
1223                 childStems = False
1224             else:
1225                 childStems = leaves * (0.1 + 0.9 * (branchL / maxbL)) * shapeRatio(leafDist, (1 - p.offset))
1226
1227         # print("n=%d, levels=%d, n'=%d, childStems=%s"%(n, levels, storeN, childStems))
1228
1229         # Determine the starting and ending radii of the stem using the tapering of the stem
1230         startRad = min((p.radiusPar[0] * ((branchL / p.lengthPar) ** ratioPower)) * radiusTweak[n], p.radiusPar[1])
1231         if p.offset == 1:
1232             startRad = p.radiusPar[1]
1233         endRad = (startRad * (1 - taper[n])) ** ratioPower
1234         startRad = max(startRad, minRadius)
1235         endRad = max(endRad, minRadius)
1236         newPoint.radius = startRad
1237
1238         # stem curvature
1239         curveVal = curve[n] / curveRes[n]
1240         curveVar = curveV[n] / curveRes[n]
1241
1242         # curveVal = curveVal * (branchL / scaleVal)
1243
1244         # Add the new stem to list of stems to grow and define which bone it will be parented to
1245         addstem(
1246             stemSpline(
1247                 newSpline, curveVal, curveVar, attractUp[n],
1248                 0, curveRes[n], branchL / curveRes[n], childStems,
1249                 startRad, endRad, len(cu.splines) - 1, 0, p.quat
1250                 )
1251             )
1252
1253         bone = roundBone(p.parBone, boneStep[n - 1])
1254         if p.offset == 1:
1255             isend = True
1256         else:
1257             isend = False
1258         addsplinetobone((bone, isend))
1259
1260
1261 def perform_pruning(baseSize, baseSplits, childP, cu, currentMax, currentMin, currentScale, curve,
1262                     curveBack, curveRes, deleteSpline, forceSprout, handles, n, oldMax, originalSplineToBone,
1263                     originalCo, originalCurv, originalCurvV, originalHandleL, originalHandleR, originalLength,
1264                     originalSeg, prune, prunePowerHigh, prunePowerLow, pruneRatio, pruneWidth, pruneBase,
1265                     pruneWidthPeak, randState, ratio, scaleVal, segSplits, splineToBone, splitAngle, splitAngleV,
1266                     st, startPrune, branchDist, length, splitByLen, closeTip, nrings, splitBias, splitHeight,
1267                     attractOut, rMode, lengthV, taperCrown, boneStep, rotate, rotateV):
1268     while startPrune and ((currentMax - currentMin) > 0.005):
1269         setstate(randState)
1270
1271         # If the search will halt after this iteration, then set the adjustment of stem
1272         # length to take into account the pruning ratio
1273         if (currentMax - currentMin) < 0.01:
1274             currentScale = (currentScale - 1) * pruneRatio + 1
1275             startPrune = False
1276             forceSprout = True
1277         # Change the segment length of the stem by applying some scaling
1278         st.segL = originalLength * currentScale
1279         # To prevent millions of splines being created we delete any old ones and
1280         # replace them with only their first points to begin the spline again
1281         if deleteSpline:
1282             for x in splineList:
1283                 cu.splines.remove(x.spline)
1284             newSpline = cu.splines.new('BEZIER')
1285             newPoint = newSpline.bezier_points[-1]
1286             newPoint.co = originalCo
1287             newPoint.handle_right = originalHandleR
1288             newPoint.handle_left = originalHandleL
1289             (newPoint.handle_left_type, newPoint.handle_right_type) = ('VECTOR', 'VECTOR')
1290             st.spline = newSpline
1291             st.curv = originalCurv
1292             st.curvV = originalCurvV
1293             st.seg = originalSeg
1294             st.p = newPoint
1295             newPoint.radius = st.radS
1296             splineToBone = originalSplineToBone
1297
1298         # Initialise the spline list for those contained in the current level of branching
1299         splineList = [st]
1300
1301         # split length variation
1302         stemsegL = splineList[0].segL  # initial segment length used for variation
1303         splineList[0].segL = stemsegL * uniform(1 - lengthV[n], 1 + lengthV[n])  # variation for first stem
1304
1305         # For each of the segments of the stem which must be grown we have to add to each spline in splineList
1306         for k in range(curveRes[n]):
1307             # Make a copy of the current list to avoid continually adding to the list we're iterating over
1308             tempList = splineList[:]
1309             # print('Leng: ', len(tempList))
1310
1311             # for curve variation
1312             if curveRes[n] > 1:
1313                 kp = (k / (curveRes[n] - 1))  # * 2
1314             else:
1315                 kp = 1.0
1316
1317             # split bias
1318             splitValue = segSplits[n]
1319             if n == 0:
1320                 splitValue = ((2 * splitBias) * (kp - .5) + 1) * splitValue
1321                 splitValue = max(splitValue, 0.0)
1322
1323             # For each of the splines in this list set the number of splits and then grow it
1324             for spl in tempList:
1325                 # adjust numSplit
1326                 lastsplit = getattr(spl, 'splitlast', 0)
1327                 splitVal = splitValue
1328                 if lastsplit == 0:
1329                     splitVal = splitValue * 1.33
1330                 elif lastsplit == 1:
1331                     splitVal = splitValue * splitValue
1332
1333                 if k == 0:
1334                     numSplit = 0
1335                 elif (n == 0) and (k < ((curveRes[n] - 1) * splitHeight)) and (k != 1):
1336                     numSplit = 0
1337                 elif (k == 1) and (n == 0):
1338                     numSplit = baseSplits
1339                 # always split at splitHeight
1340                 elif (n == 0) and (k == int((curveRes[n] - 1) * splitHeight) + 1) and (splitVal > 0):
1341                     numSplit = 1
1342                 else:
1343                     if (n >= 1) and splitByLen:
1344                         L = ((spl.segL * curveRes[n]) / scaleVal)
1345                         lf = 1
1346                         for l in length[:n + 1]:
1347                             lf *= l
1348                         L = L / lf
1349                         numSplit = splits2(splitVal * L)
1350                     else:
1351                         numSplit = splits2(splitVal)
1352
1353                 if (k == int(curveRes[n] / 2 + 0.5)) and (curveBack[n] != 0):
1354                     spl.curv += 2 * (curveBack[n] / curveRes[n])  # was -4 *
1355
1356                 growSpline(
1357                         n, spl, numSplit, splitAngle[n], splitAngleV[n], splineList,
1358                         handles, splineToBone, closeTip, kp, splitHeight, attractOut[n],
1359                         stemsegL, lengthV[n], taperCrown, boneStep, rotate, rotateV
1360                         )
1361
1362         # If pruning is enabled then we must to the check to see if the end of the spline is within the evelope
1363         if prune:
1364             # Check each endpoint to see if it is inside
1365             for s in splineList:
1366                 coordMag = (s.spline.bezier_points[-1].co.xy).length
1367                 ratio = (scaleVal - s.spline.bezier_points[-1].co.z) / (scaleVal * max(1 - pruneBase, 1e-6))
1368                 # Don't think this if part is needed
1369                 if (n == 0) and (s.spline.bezier_points[-1].co.z < pruneBase * scaleVal):
1370                     insideBool = True  # Init to avoid UnboundLocalError later
1371                 else:
1372                     insideBool = (
1373                     (coordMag / scaleVal) < pruneWidth * shapeRatio(9, ratio, pruneWidthPeak, prunePowerHigh,
1374                                                                     prunePowerLow))
1375                 # If the point is not inside then we adjust the scale and current search bounds
1376                 if not insideBool:
1377                     oldMax = currentMax
1378                     currentMax = currentScale
1379                     currentScale = 0.5 * (currentMax + currentMin)
1380                     break
1381             # If the scale is the original size and the point is inside then
1382             # we need to make sure it won't be pruned or extended to the edge of the envelope
1383             if insideBool and (currentScale != 1):
1384                 currentMin = currentScale
1385                 currentMax = oldMax
1386                 currentScale = 0.5 * (currentMax + currentMin)
1387             if insideBool and ((currentMax - currentMin) == 1):
1388                 currentMin = 1
1389
1390         # If the search will halt on the next iteration then we need
1391         # to make sure we sprout child points to grow the next splines or leaves
1392         if (((currentMax - currentMin) < 0.005) or not prune) or forceSprout:
1393             if (n == 0) and (rMode != "original"):
1394                 tVals = findChildPoints2(splineList, st.children)
1395             else:
1396                 tVals = findChildPoints(splineList, st.children)
1397             # print("debug tvals[%d] , splineList[%d], %s" % ( len(tVals), len(splineList), st.children))
1398             # If leaves is -ve then we need to make sure the only point which sprouts is the end of the spline
1399             if not st.children:
1400                 tVals = [1.0]
1401             # remove some of the points because of baseSize
1402             trimNum = int(baseSize * (len(tVals) + 1))
1403             tVals = tVals[trimNum:]
1404
1405             # grow branches in rings
1406             if (n == 0) and (nrings > 0):
1407                 # tVals = [(floor(t * nrings)) / nrings for t in tVals[:-1]]
1408                 tVals = [(floor(t * nrings) / nrings) * uniform(.995, 1.005) for t in tVals[:-1]]
1409                 tVals.append(1)
1410                 tVals = [t for t in tVals if t > baseSize]
1411
1412             # branch distribution
1413             if n == 0:
1414                 tVals = [((t - baseSize) / (1 - baseSize)) for t in tVals]
1415                 if branchDist < 1.0:
1416                     tVals = [t ** (1 / branchDist) for t in tVals]
1417                 else:
1418                     tVals = [1 - (1 - t) ** branchDist for t in tVals]
1419                 tVals = [t * (1 - baseSize) + baseSize for t in tVals]
1420
1421             # For all the splines, we interpolate them and add the new points to the list of child points
1422             maxOffset = max([s.offsetLen + (len(s.spline.bezier_points) - 1) * s.segL for s in splineList])
1423             for s in splineList:
1424                 # print(str(n)+'level: ', s.segMax*s.segL)
1425                 childP.extend(interpStem(s, tVals, s.segMax * s.segL, s.radS, maxOffset, baseSize))
1426
1427         # Force the splines to be deleted
1428         deleteSpline = True
1429         # If pruning isn't enabled then make sure it doesn't loop
1430         if not prune:
1431             startPrune = False
1432     return ratio, splineToBone
1433
1434
1435 # calculate taper automatically
1436 def findtaper(length, taper, shape, shapeS, levels, customShape):
1437     taperS = []
1438     for i, t in enumerate(length):
1439         if i == 0:
1440             shp = 1.0
1441         elif i == 1:
1442             shp = shapeRatio(shape, 0, custom=customShape)
1443         else:
1444             shp = shapeRatio(shapeS, 0)
1445         t = t * shp
1446         taperS.append(t)
1447
1448     taperP = []
1449     for i, t in enumerate(taperS):
1450         pm = 1
1451         for x in range(i + 1):
1452             pm *= taperS[x]
1453         taperP.append(pm)
1454
1455     taperR = []
1456     for i, t in enumerate(taperP):
1457         t = sum(taperP[i:levels])
1458         taperR.append(t)
1459
1460     taperT = []
1461     for i, t in enumerate(taperR):
1462         try:
1463             t = taperP[i] / taperR[i]
1464         except ZeroDivisionError:
1465             t = 1.0
1466         taperT.append(t)
1467
1468     taperT = [t * taper[i] for i, t in enumerate(taperT)]
1469
1470     return taperT
1471
1472
1473 def addTree(props):
1474     global splitError
1475     # startTime = time.time()
1476     # Set the seed for repeatable results
1477     seed(props.seed)
1478
1479     # Set all other variables
1480     levels = props.levels
1481     length = props.length
1482     lengthV = props.lengthV
1483     taperCrown = props.taperCrown
1484     branches = props.branches
1485     curveRes = props.curveRes
1486     curve = toRad(props.curve)
1487     curveV = toRad(props.curveV)
1488     curveBack = toRad(props.curveBack)
1489     baseSplits = props.baseSplits
1490     segSplits = props.segSplits
1491     splitByLen = props.splitByLen
1492     rMode = props.rMode
1493     splitAngle = toRad(props.splitAngle)
1494     splitAngleV = toRad(props.splitAngleV)
1495     scale = props.scale
1496     scaleV = props.scaleV
1497     attractUp = props.attractUp
1498     attractOut = props.attractOut
1499     shape = int(props.shape)
1500     shapeS = int(props.shapeS)
1501     customShape = props.customShape
1502     branchDist = props.branchDist
1503     nrings = props.nrings
1504     baseSize = props.baseSize
1505     baseSize_s = props.baseSize_s
1506     splitHeight = props.splitHeight
1507     splitBias = props.splitBias
1508     ratio = props.ratio
1509     minRadius = props.minRadius
1510     closeTip = props.closeTip
1511     rootFlare = props.rootFlare
1512     autoTaper = props.autoTaper
1513     taper = props.taper
1514     radiusTweak = props.radiusTweak
1515     ratioPower = props.ratioPower
1516     downAngle = toRad(props.downAngle)
1517     downAngleV = toRad(props.downAngleV)
1518     rotate = toRad(props.rotate)
1519     rotateV = toRad(props.rotateV)
1520     scale0 = props.scale0
1521     scaleV0 = props.scaleV0
1522     prune = props.prune
1523     pruneWidth = props.pruneWidth
1524     pruneBase = props.pruneBase
1525     pruneWidthPeak = props.pruneWidthPeak
1526     prunePowerLow = props.prunePowerLow
1527     prunePowerHigh = props.prunePowerHigh
1528     pruneRatio = props.pruneRatio
1529     leafDownAngle = radians(props.leafDownAngle)
1530     leafDownAngleV = radians(props.leafDownAngleV)
1531     leafRotate = radians(props.leafRotate)
1532     leafRotateV = radians(props.leafRotateV)
1533     leafScale = props.leafScale
1534     leafScaleX = props.leafScaleX
1535     leafScaleT = props.leafScaleT
1536     leafScaleV = props.leafScaleV
1537     leafShape = props.leafShape
1538     leafDupliObj = props.leafDupliObj
1539     bend = props.bend
1540     leafangle = props.leafangle
1541     horzLeaves = props.horzLeaves
1542     leafDist = int(props.leafDist)
1543     bevelRes = props.bevelRes
1544     resU = props.resU
1545
1546     useArm = props.useArm
1547     previewArm = props.previewArm
1548     armAnim = props.armAnim
1549     leafAnim = props.leafAnim
1550     frameRate = props.frameRate
1551     loopFrames = props.loopFrames
1552
1553     # windSpeed = props.windSpeed
1554     # windGust = props.windGust
1555
1556     wind = props.wind
1557     gust = props.gust
1558     gustF = props.gustF
1559
1560     af1 = props.af1
1561     af2 = props.af2
1562     af3 = props.af3
1563
1564     makeMesh = props.makeMesh
1565     armLevels = props.armLevels
1566     boneStep = props.boneStep
1567
1568     useOldDownAngle = props.useOldDownAngle
1569     useParentAngle = props.useParentAngle
1570
1571     if not makeMesh:
1572         boneStep = [1, 1, 1, 1]
1573
1574     # taper
1575     if autoTaper:
1576         taper = findtaper(length, taper, shape, shapeS, levels, customShape)
1577         # pLevels = branches[0]
1578         # taper = findtaper(length, taper, shape, shapeS, pLevels, customShape)
1579
1580     leafObj = None
1581
1582     # Some effects can be turned ON and OFF, the necessary variables are changed here
1583     if not props.bevel:
1584         bevelDepth = 0.0
1585     else:
1586         bevelDepth = 1.0
1587
1588     if not props.showLeaves:
1589         leaves = 0
1590     else:
1591         leaves = props.leaves
1592
1593     if props.handleType == '0':
1594         handles = 'AUTO'
1595     else:
1596         handles = 'VECTOR'
1597
1598     for ob in bpy.data.objects:
1599         ob.select_set(state=False)
1600
1601     # Initialise the tree object and curve and adjust the settings
1602     cu = bpy.data.curves.new('tree', 'CURVE')
1603     treeOb = bpy.data.objects.new('tree', cu)
1604     bpy.context.scene.collection.objects.link(treeOb)
1605
1606     # treeOb.location=bpy.context.scene.cursor.location attractUp
1607
1608     cu.dimensions = '3D'
1609     cu.fill_mode = 'FULL'
1610     cu.bevel_depth = bevelDepth
1611     cu.bevel_resolution = bevelRes
1612     cu.use_uv_as_generated = True
1613
1614     # Fix the scale of the tree now
1615     scaleVal = scale + uniform(-scaleV, scaleV)
1616     scaleVal += copysign(1e-6, scaleVal)  # Move away from zero to avoid div by zero
1617
1618     pruneBase = min(pruneBase, baseSize)
1619     # If pruning is turned on we need to draw the pruning envelope
1620     if prune:
1621         enHandle = 'VECTOR'
1622         enNum = 128
1623         enCu = bpy.data.curves.new('envelope', 'CURVE')
1624         enOb = bpy.data.objects.new('envelope', enCu)
1625         enOb.parent = treeOb
1626         bpy.context.scene.collection.objects.link(enOb)
1627         newSpline = enCu.splines.new('BEZIER')
1628         newPoint = newSpline.bezier_points[-1]
1629         newPoint.co = Vector((0, 0, scaleVal))
1630         (newPoint.handle_right_type, newPoint.handle_left_type) = (enHandle, enHandle)
1631         # Set the coordinates by varying the z value, envelope will be aligned to the x-axis
1632         for c in range(enNum):
1633             newSpline.bezier_points.add(1)
1634             newPoint = newSpline.bezier_points[-1]
1635             ratioVal = (c + 1) / (enNum)
1636             zVal = scaleVal - scaleVal * (1 - pruneBase) * ratioVal
1637             newPoint.co = Vector(
1638                             (
1639                             scaleVal * pruneWidth *
1640                             shapeRatio(9, ratioVal, pruneWidthPeak, prunePowerHigh, prunePowerLow),
1641                             0, zVal
1642                             )
1643                         )
1644             (newPoint.handle_right_type, newPoint.handle_left_type) = (enHandle, enHandle)
1645         newSpline = enCu.splines.new('BEZIER')
1646         newPoint = newSpline.bezier_points[-1]
1647         newPoint.co = Vector((0, 0, scaleVal))
1648         (newPoint.handle_right_type, newPoint.handle_left_type) = (enHandle, enHandle)
1649         # Create a second envelope but this time on the y-axis
1650         for c in range(enNum):
1651             newSpline.bezier_points.add(1)
1652             newPoint = newSpline.bezier_points[-1]
1653             ratioVal = (c + 1) / (enNum)
1654             zVal = scaleVal - scaleVal * (1 - pruneBase) * ratioVal
1655             newPoint.co = Vector(
1656                             (
1657                             0, scaleVal * pruneWidth *
1658                             shapeRatio(9, ratioVal, pruneWidthPeak, prunePowerHigh, prunePowerLow),
1659                             zVal
1660                             )
1661                         )
1662             (newPoint.handle_right_type, newPoint.handle_left_type) = (enHandle, enHandle)
1663
1664     childP = []
1665     stemList = []
1666
1667     levelCount = []
1668     splineToBone = deque([''])
1669     addsplinetobone = splineToBone.append
1670
1671     # Each of the levels needed by the user we grow all the splines
1672     for n in range(levels):
1673         storeN = n
1674         stemList = deque()
1675         addstem = stemList.append
1676         # If n is used as an index to access parameters for the tree
1677         # it must be at most 3 or it will reference outside the array index
1678         n = min(3, n)
1679         splitError = 0.0
1680
1681         # closeTip only on last level
1682         closeTipp = all([(n == levels - 1), closeTip])
1683
1684         # If this is the first level of growth (the trunk) then we need some special work to begin the tree
1685         if n == 0:
1686             kickstart_trunk(addstem, levels, leaves, branches, cu, curve, curveRes,
1687                             curveV, attractUp, length, lengthV, ratio, ratioPower, resU,
1688                             scale0, scaleV0, scaleVal, taper, minRadius, rootFlare)
1689         # If this isn't the trunk then we may have multiple stem to initialise
1690         else:
1691             # For each of the points defined in the list of stem starting points we need to grow a stem.
1692             fabricate_stems(addsplinetobone, addstem, baseSize, branches, childP, cu, curve, curveBack,
1693                             curveRes, curveV, attractUp, downAngle, downAngleV, leafDist, leaves, length, lengthV,
1694                             levels, n, ratioPower, resU, rotate, rotateV, scaleVal, shape, storeN,
1695                             taper, shapeS, minRadius, radiusTweak, customShape, rMode, segSplits,
1696                             useOldDownAngle, useParentAngle, boneStep)
1697
1698         # change base size for each level
1699         if n > 0:
1700             baseSize *= baseSize_s  # decrease at each level
1701         if (n == levels - 1):
1702             baseSize = 0
1703
1704         childP = []
1705         # Now grow each of the stems in the list of those to be extended
1706         for st in stemList:
1707             # When using pruning, we need to ensure that the random effects
1708             # will be the same for each iteration to make sure the problem is linear
1709             randState = getstate()
1710             startPrune = True
1711             lengthTest = 0.0
1712             # Store all the original values for the stem to make sure
1713             # we have access after it has been modified by pruning
1714             originalLength = st.segL
1715             originalCurv = st.curv
1716             originalCurvV = st.curvV
1717             originalSeg = st.seg
1718             originalHandleR = st.p.handle_right.copy()
1719             originalHandleL = st.p.handle_left.copy()
1720             originalCo = st.p.co.copy()
1721             currentMax = 1.0
1722             currentMin = 0.0
1723             currentScale = 1.0
1724             oldMax = 1.0
1725             deleteSpline = False
1726             originalSplineToBone = copy.copy(splineToBone)
1727             forceSprout = False
1728             # Now do the iterative pruning, this uses a binary search and halts once the difference
1729             # between upper and lower bounds of the search are less than 0.005
1730             ratio, splineToBone = perform_pruning(
1731                                         baseSize, baseSplits, childP, cu, currentMax, currentMin,
1732                                         currentScale, curve, curveBack, curveRes, deleteSpline, forceSprout,
1733                                         handles, n, oldMax, originalSplineToBone, originalCo, originalCurv,
1734                                         originalCurvV, originalHandleL, originalHandleR, originalLength,
1735                                         originalSeg, prune, prunePowerHigh, prunePowerLow, pruneRatio,
1736                                         pruneWidth, pruneBase, pruneWidthPeak, randState, ratio, scaleVal,
1737                                         segSplits, splineToBone, splitAngle, splitAngleV, st, startPrune,
1738                                         branchDist, length, splitByLen, closeTipp, nrings, splitBias,
1739                                         splitHeight, attractOut, rMode, lengthV, taperCrown, boneStep,
1740                                         rotate, rotateV
1741                                         )
1742
1743         levelCount.append(len(cu.splines))
1744
1745     # If we need to add leaves, we do it here
1746     leafVerts = []
1747     leafFaces = []
1748     leafNormals = []
1749
1750     leafMesh = None  # in case we aren't creating leaves, we'll still have the variable
1751
1752     leafP = []
1753     if leaves:
1754         oldRot = 0.0
1755         n = min(3, n + 1)
1756         # For each of the child points we add leaves
1757         for cp in childP:
1758             # If the special flag is set then we need to add several leaves at the same location
1759             if leaves < 0:
1760                 oldRot = -leafRotate / 2
1761                 for g in range(abs(leaves)):
1762                     (vertTemp, faceTemp, normal, oldRot) = genLeafMesh(
1763                                                                 leafScale, leafScaleX, leafScaleT,
1764                                                                 leafScaleV, cp.co, cp.quat, cp.offset,
1765                                                                 len(leafVerts), leafDownAngle, leafDownAngleV,
1766                                                                 leafRotate, leafRotateV,
1767                                                                 oldRot, bend, leaves, leafShape,
1768                                                                 leafangle, horzLeaves
1769                                                                 )
1770                     leafVerts.extend(vertTemp)
1771                     leafFaces.extend(faceTemp)
1772                     leafNormals.extend(normal)
1773                     leafP.append(cp)
1774             # Otherwise just add the leaves like splines
1775             else:
1776                 (vertTemp, faceTemp, normal, oldRot) = genLeafMesh(
1777                                                             leafScale, leafScaleX, leafScaleT, leafScaleV,
1778                                                             cp.co, cp.quat, cp.offset, len(leafVerts),
1779                                                             leafDownAngle, leafDownAngleV, leafRotate,
1780                                                             leafRotateV, oldRot, bend, leaves, leafShape,
1781                                                             leafangle, horzLeaves
1782                                                             )
1783                 leafVerts.extend(vertTemp)
1784                 leafFaces.extend(faceTemp)
1785                 leafNormals.extend(normal)
1786                 leafP.append(cp)
1787
1788         # Create the leaf mesh and object, add geometry using from_pydata,
1789         # edges are currently added by validating the mesh which isn't great
1790         leafMesh = bpy.data.meshes.new('leaves')
1791         leafObj = bpy.data.objects.new('leaves', leafMesh)
1792         bpy.context.scene.collection.objects.link(leafObj)
1793         leafObj.parent = treeOb
1794         leafMesh.from_pydata(leafVerts, (), leafFaces)
1795
1796         # set vertex normals for dupliVerts
1797         if leafShape == 'dVert':
1798             leafMesh.vertices.foreach_set('normal', leafNormals)
1799
1800         # enable duplication
1801         if leafShape == 'dFace':
1802             leafObj.instance_type = "FACES"
1803             leafObj.use_instance_faces_scale = True
1804             leafObj.instance_faces_scale = 10.0
1805             try:
1806                 if leafDupliObj not in "NONE":
1807                     bpy.data.objects[leafDupliObj].parent = leafObj
1808             except KeyError:
1809                 pass
1810         elif leafShape == 'dVert':
1811             leafObj.instance_type = "VERTS"
1812             leafObj.use_instance_vertices_rotation = True
1813             try:
1814                 if leafDupliObj not in "NONE":
1815                     bpy.data.objects[leafDupliObj].parent = leafObj
1816             except KeyError:
1817                 pass
1818
1819         # add leaf UVs
1820         if leafShape == 'rect':
1821             leafMesh.uv_layers.new(name='leafUV')
1822             uvlayer = leafMesh.uv_layers.active.data
1823
1824             u1 = .5 * (1 - leafScaleX)
1825             u2 = 1 - u1
1826
1827             for i in range(0, len(leafFaces)):
1828                 uvlayer[i * 4 + 0].uv = Vector((u2, 0))
1829                 uvlayer[i * 4 + 1].uv = Vector((u2, 1))
1830                 uvlayer[i * 4 + 2].uv = Vector((u1, 1))
1831                 uvlayer[i * 4 + 3].uv = Vector((u1, 0))
1832
1833         elif leafShape == 'hex':
1834             leafMesh.uv_layers.new(name='leafUV')
1835             uvlayer = leafMesh.uv_layers.active.data
1836
1837             u1 = .5 * (1 - leafScaleX)
1838             u2 = 1 - u1
1839
1840             for i in range(0, int(len(leafFaces) / 2)):
1841                 uvlayer[i * 8 + 0].uv = Vector((.5, 0))
1842                 uvlayer[i * 8 + 1].uv = Vector((u1, 1 / 3))
1843                 uvlayer[i * 8 + 2].uv = Vector((u1, 2 / 3))
1844                 uvlayer[i * 8 + 3].uv = Vector((.5, 1))
1845
1846                 uvlayer[i * 8 + 4].uv = Vector((.5, 0))
1847                 uvlayer[i * 8 + 5].uv = Vector((.5, 1))
1848                 uvlayer[i * 8 + 6].uv = Vector((u2, 2 / 3))
1849                 uvlayer[i * 8 + 7].uv = Vector((u2, 1 / 3))
1850
1851         leafMesh.validate()
1852
1853     leafVertSize = {'hex': 6, 'rect': 4, 'dFace': 4, 'dVert': 1}[leafShape]
1854
1855     armLevels = min(armLevels, levels)
1856     armLevels -= 1
1857
1858     # unpack vars from splineToBone
1859     splineToBone1 = splineToBone
1860     splineToBone = [s[0] if len(s) > 1 else s for s in splineToBone1]
1861     isend = [s[1] if len(s) > 1 else False for s in splineToBone1]
1862     issplit = [s[2] if len(s) > 2 else False for s in splineToBone1]
1863     splitPidx = [s[3] if len(s) > 2 else 0 for s in splineToBone1]
1864
1865     # If we need an armature we add it
1866     if useArm:
1867         # Create the armature and objects
1868         create_armature(
1869                     armAnim, leafP, cu, frameRate, leafMesh, leafObj, leafVertSize,
1870                     leaves, levelCount, splineToBone, treeOb, wind, gust, gustF, af1,
1871                     af2, af3, leafAnim, loopFrames, previewArm, armLevels, makeMesh, boneStep
1872                     )
1873
1874     # print(time.time()-startTime)
1875
1876     # mesh branches
1877     if makeMesh:
1878         t1 = time.time()
1879
1880         treeMesh = bpy.data.meshes.new('treemesh')
1881         treeObj = bpy.data.objects.new('treemesh', treeMesh)
1882         bpy.context.scene.collection.objects.link(treeObj)
1883
1884         treeVerts = []
1885         treeEdges = []
1886         root_vert = []
1887         vert_radius = []
1888         vertexGroups = OrderedDict()
1889         lastVerts = []
1890
1891         for i, curve in enumerate(cu.splines):
1892             points = curve.bezier_points
1893
1894             # find branching level
1895             level = 0
1896             for l, c in enumerate(levelCount):
1897                 if i < c:
1898                     level = l
1899                     break
1900             level = min(level, 3)
1901
1902             step = boneStep[level]
1903             vindex = len(treeVerts)
1904
1905             p1 = points[0]
1906
1907             # add extra vertex for splits
1908             if issplit[i]:
1909                 pb = int(splineToBone[i][4:-4])
1910                 pn = splitPidx[i]  # int(splineToBone[i][-3:])
1911                 p_1 = cu.splines[pb].bezier_points[pn]
1912                 p_2 = cu.splines[pb].bezier_points[pn + 1]
1913                 p = evalBez(p_1.co, p_1.handle_right, p_2.handle_left, p_2.co, 1 - 1 / (resU + 1))
1914                 treeVerts.append(p)
1915
1916                 root_vert.append(False)
1917                 vert_radius.append((p1.radius * .75, p1.radius * .75))
1918                 treeEdges.append([vindex, vindex + 1])
1919                 vindex += 1
1920
1921             if isend[i]:
1922                 parent = lastVerts[int(splineToBone[i][4:-4])]
1923                 vindex -= 1
1924             else:
1925                 # add first point
1926                 treeVerts.append(p1.co)
1927                 root_vert.append(True)
1928                 vert_radius.append((p1.radius, p1.radius))
1929             """
1930             # add extra vertex for splits
1931             if issplit[i]:
1932                 p2 = points[1]
1933                 p = evalBez(p1.co, p1.handle_right, p2.handle_left, p2.co, .001)
1934                 treeVerts.append(p)
1935                 root_vert.append(False)
1936                 vert_radius.append((p1.radius, p1.radius)) #(p1.radius * .95, p1.radius * .95)
1937                 treeEdges.append([vindex,vindex+1])
1938                 vindex += 1
1939             """
1940             # dont make vertex group if above armLevels
1941             if (i >= levelCount[armLevels]):
1942                 idx = i
1943                 groupName = splineToBone[idx]
1944                 g = True
1945                 while groupName not in vertexGroups:
1946                     # find parent bone of parent bone
1947                     b = splineToBone[idx]
1948                     idx = int(b[4:-4])
1949                     groupName = splineToBone[idx]
1950             else:
1951                 g = False
1952
1953             for n, p2 in enumerate(points[1:]):
1954                 if not g:
1955                     groupName = 'bone' + (str(i)).rjust(3, '0') + '.' + (str(n)).rjust(3, '0')
1956                     groupName = roundBone(groupName, step)
1957                     if groupName not in vertexGroups:
1958                         vertexGroups[groupName] = []
1959
1960                 # parent first vert in split to parent branch bone
1961                 if issplit[i] and n == 0:
1962                     if g:
1963                         vertexGroups[groupName].append(vindex - 1)
1964                     else:
1965                         vertexGroups[splineToBone[i]].append(vindex - 1)
1966
1967                 for f in range(1, resU + 1):
1968                     pos = f / resU
1969                     p = evalBez(p1.co, p1.handle_right, p2.handle_left, p2.co, pos)
1970                     radius = p1.radius + (p2.radius - p1.radius) * pos
1971
1972                     treeVerts.append(p)
1973                     root_vert.append(False)
1974                     vert_radius.append((radius, radius))
1975
1976                     if (isend[i]) and (n == 0) and (f == 1):
1977                         edge = [parent, n * resU + f + vindex]
1978                     else:
1979                         edge = [n * resU + f + vindex - 1, n * resU + f + vindex]
1980                         # add vert to group
1981                         vertexGroups[groupName].append(n * resU + f + vindex - 1)
1982                     treeEdges.append(edge)
1983
1984                 vertexGroups[groupName].append(n * resU + resU + vindex)
1985
1986                 p1 = p2
1987
1988             lastVerts.append(len(treeVerts) - 1)
1989
1990         treeMesh.from_pydata(treeVerts, treeEdges, ())
1991
1992         for group in vertexGroups:
1993             treeObj.vertex_groups.new(name=group)
1994             treeObj.vertex_groups[group].add(vertexGroups[group], 1.0, 'ADD')
1995
1996         # add armature
1997         if useArm:
1998             armMod = treeObj.modifiers.new('windSway', 'ARMATURE')
1999             if previewArm:
2000                 bpy.data.objects['treeArm'].hide_viewport = True
2001                 bpy.data.armatures['tree'].display_type = 'STICK'
2002             armMod.object = bpy.data.objects['treeArm']
2003             armMod.use_bone_envelopes = False
2004             armMod.use_vertex_groups = True
2005             treeObj.parent = bpy.data.objects['treeArm']
2006
2007         # add skin modifier and set data
2008         skinMod = treeObj.modifiers.new('Skin', 'SKIN')
2009         skinMod.use_smooth_shade = True
2010         if previewArm:
2011             skinMod.show_viewport = False
2012         skindata = treeObj.data.skin_vertices[0].data
2013         for i, radius in enumerate(vert_radius):
2014             skindata[i].radius = radius
2015             skindata[i].use_root = root_vert[i]
2016
2017         print("mesh time", time.time() - t1)