text: # The objective is to build a tree object whose all branches are connected. In other words, the tree is manifold.
  # To do so, the tree is generated as a succession of modules that can represent splits, stems and so on.
  # The modules are objects from the class Module.
  # Each module has a resolution, except the modules that make the junction between two different levels of subdivision
  
  
  
  
  import bpy, bmesh
  import numpy as np
  from mathutils import Vector, Matrix
  from math import pi, sqrt, cos, sin
  from .bridge import bridge
  from random import random
  
  
  def square(size):
      """Returns a list of 4 vectors arranged in a square of specified size"""
      return [Vector((-1, -1, 0))*size, Vector((1, -1, 0))*size, Vector((1, 1, 0))*size, Vector((-1, 1, 0))*size]
  
  
  def octagon(size):
      result = []
      for i in range(8):
          angle = pi*i/4
          result.append(Vector((cos(angle), sin(angle), 0)) * size)
      return result
  
  
  def directions_to_spin(direction, secondary_direction):
      direction_rotation = Vector((0,0,1)).rotation_difference(direction).to_matrix()
      secondary_direction = secondary_direction * direction_rotation
      spin = - secondary_direction.xy.angle_signed(Vector((-1, 0)))
      return spin
  
  
  def get_direction(primary_direction, angle, spin):
      rot_1 = Matrix.Rotation(angle, 4, 'Y')
      rot_2 = primary_direction.rotation_difference(Vector((0,0,1))).to_matrix()
      rot_3 = Matrix.Rotation(spin, 4, 'Z')
      direction = ((rot_1 * Vector((0, 0, 1))) * rot_3) * rot_2
      return direction
  
  
  def average_vector(vectors):
      """returns the average vector of a list of vectors"""
      v = sum(vectors, Vector())
      v /= len(vectors)
      return v
  
  
  def find_verts_number_rec(module):
      if module is None:
          return 0
      if module.type == 'root' or module.type == 'branch':
          return 4 + find_verts_number_rec(module.head_module_1)
      if module.type == 'split':
          return 8 + find_verts_number_rec(module.head_module_1) + find_verts_number_rec(module.head_module_2)
  
  
  def find_faces_number_rec(module):
      if module is None:
          return 1
      if module.type == 'root':
          return find_faces_number_rec(module.head_module_1)
      if module.type == 'branch':
          return 4 + find_faces_number_rec(module.head_module_1)
      if module.type == 'split':
          return 7 + find_faces_number_rec(module.head_module_1) + find_faces_number_rec(module.head_module_2)
  
  
  def draw_module(root, resolution_levels, twig=False):
      max_radius = root.base_radius
      root.resolution = resolution_levels
      apply_resolution_rec(root.head_module_1, resolution_levels, max_radius, root)
  
      verts = [[] for i in range(resolution_levels+1)]
      faces = [[] for i in range(resolution_levels+1)]
      uvs = [[] for i in range(resolution_levels+1)]
      v_groups = [[] for i in range(resolution_levels+1)]
  
      root.build()
      verts[-1] = [v for v in root.verts]
      extremities = [root]
      while len(extremities) > 0:
          new_extremities = []
          for module in extremities:
              for head in range(module.head_number):
                  new_module = module.head_module_1 if head == 0 else module.head_module_2
                  if new_module is not None:
                      resolution = new_module.resolution
                      curr_verts_number = len(verts[resolution])
                      module.link(new_module, head, curr_verts_number)
                      verts[resolution].extend(new_module.verts)
                      faces[resolution].extend(new_module.faces)
                      uvs[resolution].extend(new_module.uvs)
                      v_groups[resolution].append((list(range(curr_verts_number, curr_verts_number+len(new_module.verts))), new_module.base_radius / max_radius))
                      new_extremities.append(new_module)
  
          extremities = new_extremities
  
      objects = []
  
      for i in range(resolution_levels+1):
          bpy.ops.object.select_all(action='DESELECT')
          name = "twig" if twig else "tree"
          mesh = bpy.data.meshes.new(name)
          bm = bmesh.new()
          bm.from_mesh(mesh)
  
          for v in verts[i]:
              bm.verts.new(v)
          bm.verts.ensure_lookup_table()
  
          for f in faces[i]:
              try:
                  bm.faces.new([bm.verts[j] for j in f])
              except:
                  print(f)
  
          bm.faces.ensure_lookup_table()
  
          bm.loops.layers.uv.new()
          uv_layer = bm.loops.layers.uv.active
          for index, face in enumerate(bm.faces):
              for j, loop in enumerate(face.loops):
                  loop[uv_layer].uv = uvs[i][index][j]
  
          bm.to_mesh(mesh)
          bm.free()
          obj = bpy.data.objects.new(name, mesh)
          obj.location = bpy.context.scene.cursor_location
          vg = obj.vertex_groups.new("radius")
          for weights in v_groups[i]:
              vg.add(weights[0], weights[1], 'REPLACE')
          bpy.context.scene.objects.link(obj)
          bpy.context.scene.objects.active = obj
          if not twig:
              obj["is_tree"] = True
  
          obj["tree_type"] = "object"
          obj.select = True
          bpy.ops.object.mode_set(mode='EDIT')
          bpy.ops.mesh.select_all(action='SELECT')
          bpy.ops.mesh.normals_make_consistent(inside=False)
          bpy.ops.object.mode_set(mode='OBJECT')
          bpy.ops.object.subdivision_set(level=i)
          print(i)
          objects.append(obj)
  
      for o in objects:
          o.select = True
      bpy.ops.object.convert(target='MESH')
      bpy.ops.object.join()
      bridge(bpy.context.object)
  
  
  def visualize_with_curves(root):
      curve_data = bpy.data.curves.new('Tree', type='CURVE')
      curve_data.dimensions = '3D'
      polyline = curve_data.splines.new('POLY')
      x, y, z = root.position
      polyline.points[0].co = (x, y, z, 1)
      polyline.points[0].radius = root.base_radius
      draw_curve_rec(root, polyline, curve_data)
  
      curveOB = bpy.data.objects.new('Tree', curve_data)
      curveOB.location = bpy.context.scene.cursor_location
      curve_data.bevel_depth = 1
      curve_data.bevel_resolution = 0
      curve_data.fill_mode = 'FULL'
  
      scene = bpy.context.scene
      scene.objects.link(curveOB)
      scene.objects.active = curveOB
      curveOB["is_tree"] = True
      curveOB["tree_type"] = "curve"
      curveOB.select = True
  
  
  def draw_curve_rec(module, polyline, curve_data):
      if module is not None:
          polyline.points.add(1)
          x,y,z = module.position
          polyline.points[-1].co = (x, y, z, 1)
          polyline.points[-1].radius = module.base_radius
          draw_curve_rec(module.head_module_1, polyline, curve_data)
          if module.type == 'split' and module.head_module_2 is not None:
              new_polyline = curve_data.splines.new('POLY')
              # new_polyline.points.add(1)
              new_polyline.points[0].co = (x, y, z, 1)
              new_polyline.points[0].radius = module.base_radius
              draw_curve_rec(module.head_module_2, new_polyline, curve_data)
  
  
  def roll_indexes(indexes, angle_diff):
      """ shift a list according to an angle between two squares so that the indexes are aligned"""
      shift = int(4*angle_diff/pi)%8
      shifts = np.array([0, -1, -1, -2, -2, -3, -3, 0, 0, 0])
      shift = shifts[shift]
      return np.roll(indexes, -shift)
  
  
  def apply_resolution_rec(module, resolution_levels, max_radius, parent):
      resolution = max(parent.resolution - 1, int(resolution_levels * module.base_radius / max_radius +.6))
      if module.type == 'branch' and parent.type == 'branch':
          module.resolution = resolution
          if resolution < parent.resolution:
              # transition = Transition(module.position, module.direction, module.base_radius, module.length, module.head_1_radius, resolution, module.starting_index, module.spin)
              # if parent.head_module_2 == module:
              #     parent.head_module_2 = transition
              # else:
              #     parent.head_module_1 = transition
              # transition.head_module_1 = module.head_module_1
              # transition.head_module_2 = module.head_module_2
              # module = transition
              module.direction = parent.direction
              module.draw_base = True
      else:
          module.resolution = parent.resolution
  
      if module.head_module_1 is not None:
          apply_resolution_rec(module.head_module_1, resolution_levels, max_radius, module)
  
      if module.head_module_2 is not None:
          apply_resolution_rec(module.head_module_2, resolution_levels, max_radius, module)
  
  
  class Module:
      def __init__(self, position, direction, radius, resolution, starting_index, spin):
          self.verts = np.array([])
          self.faces = np.array([])
          self.uvs = np.array([])
          self.uv_height = 0
          self.type = 'module'
          self.creator = "default"
          self.position = position
          self.direction = direction
          self.base_radius = radius
          self.resolution = resolution
          self.spin = spin
          self.starting_index = starting_index
          # self.base_pos = []
          self.head_module_1 = None
          self.head_module_2 = None
  
      def draw(self):
          mesh = bpy.data.meshes.new(self.type)
          bm = bmesh.new()
          bm.from_mesh(mesh)
          for v in self.verts:
              bm.verts.new(v)
          bm.verts.ensure_lookup_table()
  
          for f in self.faces:
              bm.faces.new([bm.verts[i] for i in f])
  
          bm.to_mesh(mesh)
          bm.free()
          obj = bpy.data.objects.new(self.type, mesh)
          obj.location = Vector((0, 0, 0))
          bpy.context.scene.objects.link(obj)
          bpy.context.scene.objects.active = obj
          obj.select = True
  
      def get_faces(self, starting_index=0):
          return [tuple([i + starting_index for i in f]) for f in self.faces]
  
      def get_head_direction(self, head):
          return self.direction
  
      def get_extremities_rec(self, curr_extremities, selection):
          is_selected = selection == [] or self.creator in selection
          if self.head_module_1 is None:
              if is_selected:
                  curr_extremities.append((self, 0))
          else:
              self.head_module_1.get_extremities_rec(curr_extremities, selection)
          if self.head_module_2 is None:
              if self.type == 'split' and is_selected:
                  curr_extremities.append((self, 1))
          else:
              self.head_module_2.get_extremities_rec(curr_extremities, selection)
  
      def __repr__(self):
          if self.type == 'split':
              return str(self.type) + " " + str(self.position.to_tuple(2)) + " " + self.head_module_1.__repr__() + self.head_module_2.__repr__()
          else:
              return str(self.type) + " " + str(self.position.to_tuple(2)) + " " + self.head_module_1.__repr__()
  
  
  class Split(Module):
      def __init__(self, position=Vector(), direction=Vector(), radius=1, resolution=0, starting_index=0, spin=0, head_2_length=1, head_2_radius=.6):
          Module.__init__(self,position, direction, radius, resolution, starting_index, spin)
          self.type = 'split'
          self.head_1_radius = .99 * self.base_radius
          self.head_2_radius = head_2_radius * self.base_radius
          self.primary_angle = pi/18
          self.secondary_angle = pi/4
          self.head_1_length = self.base_radius * 3
          self.head_2_length = head_2_length
          self.head_number = 2
          self.head_1_direction = Vector()
          self.head_2_direction = Vector()
  
  
          # self.build()
          # self.draw()
  
      def get_head_indexes(self, head):
          if head == 1:
              return [self.starting_index + i for i in range(4)]
          else:
              return [self.starting_index + 4 + i for i in range(4)]
  
      def get_head_pos(self, head):
          if head==0:
              direction = self.get_head_direction(0)
              return self.position + direction * self.head_1_length
          elif head==1:
              direction = self.get_head_direction(1)
              return self.position + direction * self.head_2_length
  
      def get_head_direction(self, head):
          if head == 0:
              return get_direction(self.direction, self.primary_angle, self.spin)
          else:
              return get_direction(self.direction, self.primary_angle - self.secondary_angle, self.spin)
  
      def build(self, base_indexes=range(4)):
          radius_correction = 1 - .25**(self.resolution+1)
          self.base_radius *= radius_correction
          self.head_1_radius *= radius_correction
          self.head_2_radius *= radius_correction
  
          uv_height = self.uv_height
          v2 = square(self.head_1_radius)
          v3 = square(self.head_2_radius)
          primary_rotation = Matrix.Rotation(self.primary_angle, 4, 'Y')
          secondary_rotation = Matrix.Rotation(self.primary_angle - self.secondary_angle, 4, 'Y')
          v2 = [primary_rotation * (v + Vector((0, 0, self.head_1_length))) for v in v2]
          v3 = [secondary_rotation * (v + Vector((0, 0, self.head_2_length))) for v in v3]
          self.verts = v2 + v3
          si = self.starting_index
          i0, i1, i2, i3 = base_indexes
          faces = [(i0, si, si + 1, i1), (i1, si + 1, si + 2, i2), (i2, si + 2, si + 3, i3), (i3, si+3, si+6, si+7), (si+4, si+5, si, i0), (si+6, si+3, si, si+5), (i3, si+7, si+4, i0)]
  
          spin_rotation = Matrix.Rotation(self.spin, 4, 'Z')
          direction_rotation = self.direction.rotation_difference(Vector((0,0,1))).to_matrix()
          self.verts = [((v * spin_rotation) * direction_rotation) + self.position for v in self.verts]
          self.head_1_direction = (self.verts[-5] + self.verts[-7])/2 - self.position
          self.head_2_direction = (self.verts[-1] + self.verts[-3])/2 - self.position
          self.verts = np.asarray([i.to_tuple() for i in self.verts])
          # self.base_pos = [(v * spin_rotation) * direction_rotation + self.position for v in square(self.base_radius)]
          uvs = [[(i / 4, uv_height), (i / 4, uv_height + .1*self.head_1_length / self.head_1_radius), ((i + 1) / 4, uv_height + .1*self.head_1_length / self.head_1_radius), ((i + 1) / 4, uv_height)] for i in range(3)]
          uvs.extend([[(i / 4, uv_height), (i / 4, uv_height + .1*self.head_2_length / self.head_2_radius), ((i + 1) / 4, uv_height + .1*self.head_2_length / self.head_2_radius), ((i + 1) / 4, uv_height)] for i in range(4)])
  
          if self.head_module_1 is None:
              faces.append([i for i in range(si, si+4)])
              uvs.append([(0, 0), (0, 1), (1, 1), (1, 0)])
  
          if self.head_module_2 is None:
              faces.append([i for i in range(si+4, si + 8)])
              uvs.append([(0, 0), (0, 1), (1, 1), (1, 0)])
  
          self.faces = np.asarray(faces)
          self.uvs = np.asarray(uvs)
  
      def link(self, module, head, verts_number):
          if head ==0:
              head_indexes = np.arange(self.starting_index, self.starting_index + 4)
          else:
              head_indexes = np.arange(self.starting_index+4, self.starting_index + 8)
  
          module.starting_index = verts_number
          spin_diff = (module.spin - self.spin) % (2*pi)
          module.build(roll_indexes(head_indexes, spin_diff))
  
  
  class Branch(Module):
      def __init__(self, position=Vector(), direction=Vector, radius=1, length=1, head_radius=.95, resolution=0,
                   starting_index=0, spin=0):
          Module.__init__(self, position, direction, radius, resolution, starting_index, spin)
          self.type = "branch"
          self.length = length
          self.head_1_radius = head_radius * self.base_radius
          self.head_number = 1
          self.draw_base = False
  
          # self.build()
          # self.draw()
  
      def get_head_indexes(self, head):
          if head == 1:
              return np.arange(4) + self.starting_index
  
      def get_head_pos(self, head):
          return self.position + self.direction * self.length
  
      def build(self, base_indexes=np.arange(4)):
          radius_correction = 1 - .25**(self.resolution+1)
          self.base_radius *= radius_correction
          self.head_1_radius *= radius_correction
          uv_height = self.uv_height
          v2 = [v + Vector((0,0, self.length)) for v in square(self.head_1_radius)]
          if self.draw_base:
              v2.extend([v + Vector((0,0, 0*self.length/4)) for v in square(self.base_radius)])
          self.verts = v2
  
          i0, i1, i2, i3 = base_indexes
          si = self.starting_index
          if self.draw_base:
              i0, i1, i2, i3 = [si + 4 + i for i in range(4)]
  
          spin_rotation = Matrix.Rotation(self.spin, 4, 'Z')
          direction_rotation = self.direction.rotation_difference(Vector((0, 0, 1))).to_matrix()
          self.verts = np.asarray([((v * spin_rotation) * direction_rotation + self.position).to_tuple() for v in self.verts])
          # self.base_pos = [(v * spin_rotation) * direction_rotation + self.position for v in square(self.base_radius)]
          m = min(base_indexes)
          faces = [(i0, si, si + 1, i1), (i1, si + 1, si + 2, i2), (i2, si + 2, si + 3, i3), (i3, si + 3, si, i0)]
          uvs = [[(i/4, uv_height), (i/4, uv_height + .1*self.length/self.base_radius), ((i+1)/4, uv_height + .1*self.length/self.base_radius), ((i+1)/4, uv_height)] for i in range(4)]
  
          if self.head_module_1 is None:
              faces.append([i for i in range(si, si+4)])
              uvs.append([(0, 0), (0, 1), (1, 1), (1, 0)])
  
          self.faces = np.asarray(faces)
          self.uvs = np.asarray(uvs)
  
      def link(self, module, head, verts_number):
          # module.position = (self.verts[0] + self.verts[2])/2
          # module.base_radius = self.head_1_radius
          # module.direction = self.direction
          module.starting_index = verts_number
          # self.head_module_1 = module
          n = len(self.verts)
          head_indexes = np.arange(4) + self.starting_index
          spin_diff = module.spin - self.spin
          module.uv_height = self.uv_height + .1*self.length/self.base_radius
          module.build(roll_indexes(head_indexes, spin_diff))
          # base_verts = roll_indexes(module.base_pos, spin_diff)
          # for i in range(4):
          #     self.verts[i] = self.verts[i]*.7 + base_verts[i]*.3
  
  
  class Root(Module):
      def __init__(self, position=Vector(), direction=Vector, radius=1, resolution=0, starting_index=0, spin=0):
          Module.__init__(self, position, direction, radius, resolution, starting_index, spin)
          self.type = "root"
          self.head_number = 1
          self.head_1_radius = self.base_radius
          self.density_dict = dict()
  
      def build(self):
          spin_rotation = Matrix.Rotation(self.spin, 4, 'Z')
          direction_rotation = self.direction.rotation_difference(Vector((0, 0, 1))).to_matrix()
          self.verts = np.asarray([((v * spin_rotation) * direction_rotation + self.position).to_tuple() for v in square(self.base_radius)])
  
      def get_head_pos(self, head):
          return self.position
  
      def link(self, module, head, verts_number):
          module.starting_index = verts_number
          spin_diff = module.spin - self.spin
          head_indexes = np.arange(4)
          module.uv_height = 0
          module.build(roll_indexes(head_indexes, spin_diff))
  
  
  class Transition(Module):
      def __init__(self, position=Vector(), direction=Vector(), radius=1, length=1, head_radius=.95, resolution=0, starting_index=0, spin=0 ):
          Module.__init__(self, position, direction, radius, resolution, starting_index, spin)
          self.type = "transition"
          self.head_1_radius = head_radius
          self.length = length
          self.head_number = 1
          # print('transition')
  
      def build(self, base_indexes=None, create_base=False):
          v1 = octagon(self.base_radius)
          radius_correction = 1 - .25 ** (self.resolution + 1)
          self.base_radius *= radius_correction
          self.head_1_radius *= radius_correction
          v2 = [v + Vector((0,0, self.length)) for v in square(self.base_radius)]
          # v3 = [v * self.head_1_radius + Vector((0, 0, self.length/2)) for v in [Vector((-1, -.5, 0)), Vector((1, -.5, 0)), Vector((1, .5, 0)), Vector((-1, .5, 0))]]
  
          verts = v1 + v2
  
          faces = np.array([[0, 1, 14, 13], [1, 2, 10, 14], [2, 3, 11, 10], [3, 4, 15, 11], [4, 5, 12, 15], [5, 6, 8, 12], [6, 7, 9, 8], [7, 0, 13, 9], [13, 14, 10, 9], [15, 12, 8, 11]])
          faces = [[0, 1, 10, 9], [1,2,10], [2,3,11,10], [3,4,11], [4,5,8,11], [5,6,8], [6,7,9,8], [7,0,9]]
          # faces = np.array([[0, 1, 14, 13]])
          # self.faces = faces + self.starting_index
          self.faces = [tuple([i + self.starting_index for i in j]) for j in faces]
          # verts = [Vector(i) * self.base_radius for i in [(-0.37, 0.9, 0.0), (-0.9, 0.37, 0.0), (-0.9, -0.37, 0.0),
          #                                                 (-0.37, -0.9, 0.0), (0.37, -0.9, 0.0), (0.9, -0.37, -0.0),
          #                                                 (0.9, 0.37, 0.0), (0.37, 0.9, 0.0), (-0.9, -0.9, 1.38),
          #                                                 (0.9, -0.9, 1.38),
          #                                                 (0.9, 0.9, 1.38), (-0.9, 0.9, 1.38), (-0.93, 0.42, 0.84),
          #                                                 (-0.93, -0.42, 0.84), (0.93, -0.42, 0.84), (0.93, 0.42, 0.84)]]
          #
          # f = [tuple([i + self.starting_index for i in j]) for j in [[6, 15, 14, 5], [2, 13, 12, 1], [8, 13, 2, 3],
          #                                                            [11, 12, 13, 8], [0, 1, 12, 11], [14, 9, 4, 5],
          #                                                            [10, 9, 14, 15], [6, 7, 10, 15], [4, 9, 8, 3],
          #                                                            [0, 11, 10, 7]]]
  
          spin_rotation = Matrix.Rotation(self.spin, 4, 'Z')
          direction_rotation = self.direction.rotation_difference(Vector((0, 0, 1))).to_matrix()
          self.verts = np.asarray([((v * spin_rotation) * direction_rotation + self.position).to_tuple() for v in verts])
  
  
      def link(self, module, head, verts_number):
          # module.position = (self.verts[0] + self.verts[2])/2
          # module.base_radius = self.head_1_radius
          # module.direction = self.direction
          module.starting_index = verts_number
          # self.head_module_1 = module
          head_indexes = np.arange(4) + self.starting_index + 8
          spin_diff = module.spin - self.spin
          module.uv_height = self.uv_height + .1*self.length/self.base_radius
          module.build(roll_indexes(head_indexes, spin_diff))
          # base_verts = roll_indexes(module.base_pos, spin_diff)
          # for i in range(4):
          #     self.verts[i] = self.verts[i]*.7 + base_verts[i]*.3
  
  