import bpy
from random import randint, random
from mathutils import Vector
from math import pi, sin, cos


bpy.ops.object.select_all(action='SELECT')
bpy.ops.object.delete(use_global=False)

def new_tree():
    verts = []
    faces = []
    face_uvs = []
    cylinder_sections = 30
    bpy.ops.mesh.primitive_ico_sphere_add(subdivisions=2, radius=1)
    top_fruit = bpy.context.active_object    
    top_fruit.location = (0, 0, 0)
    # For Leaves
    # leaf_faces
    # leaf_verts

    def make_end_ring(start, end, radius, start_instead=False):
        middle_vector = end - start
        point_one = middle_vector.cross((1, 0, 0)).normalized() * radius
        point_two = middle_vector.cross(point_one).normalized() * radius
        ring_indices = []
        for s in range(cylinder_sections + 1):
            angle_a = s * 2 * pi / cylinder_sections
            start_edge_point = cos(angle_a) * point_one + sin(angle_a) * point_two + start
            end_edge_point = start_edge_point + middle_vector
            ring_indices.append(len(verts))
            if start_instead:
               verts.append(start_edge_point)
            else:
               verts.append(end_edge_point)
        return ring_indices
    
    def make_section(start_ring, end_ring):
        for i in range(cylinder_sections):
            faces.append([start_ring[i], end_ring[i], end_ring[i+1], start_ring[i+1]])
            # We assign texture coordinates right here
            # 0 and 1 indicate that we use the whole height of the texture image on this section
            # i/cylinder_sections is responsible for the x range we're using in our texture
            face_uvs.append(((i / cylinder_sections, 0), (i/cylinder_sections, 1))) 
        
    # run this AFTER the tree mesh is created
    def add_texture_coordinates(tree_mesh):
        uv_layer = tree_mesh.uv_layers.new(name="Bark")
        for mpoly, fuv in zip(tree_mesh.polygons, face_uvs):
            uv_layer.data[mpoly.loop_indices[0]].uv = fuv[0]
            uv_layer.data[mpoly.loop_indices[1]].uv = fuv[1]
            # For the next:  3, than 2
            uv_layer.data[mpoly.loop_indices[3]].uv = (fuv[0][0] + (1/cylinder_sections), fuv[0][1])
            uv_layer.data[mpoly.loop_indices[2]].uv = (fuv[1][0] + (1/cylinder_sections), fuv[1][1])
    
    # Add leaf texture coordinates in a similar way

#    def draw_branch(start, end, radius, prev_pos, subdivisions, start_ring):
#        # Tuesday's bug:  I used end instead of prev_pos
#        # Really, it was that simple!
#        section_length = (start - prev_pos).length
#        continue_vector = section_length * (start - prev_pos).normalized()
#
#        middle_of_straight_branch = (start + end) / 2
#        end_of_continue = start + continue_vector
#        half_continue = start + continue_vector/2
#        early_spot = (2 * start + middle_of_straight_branch + 2 * half_continue ) / 5
#        # I did a little re-weighting on middle_spot and late_spot to avoid abrupt jogs
#        middle_spot = (2 * start + end + end_of_continue) / 4
#        late_spot = (middle_spot + 2 * end + middle_of_straight_branch) / 4
#        early_ring = make_end_ring(start, early_spot, radius * 1.2)
#        middle_ring = make_end_ring(early_spot, middle_spot, radius * 1.1)
#        late_ring = make_end_ring(middle_spot, late_spot, radius * 1.05)
#        end_ring = make_end_ring(late_spot, end, radius)
#        make_section(start_ring, early_ring);
#        # draw_branch(start, early_spot, radius * 1.2, prev_pos, subdivisions - 1, start_ring)
#        make_section(early_ring, middle_ring);
#        make_section(middle_ring, late_ring)
#        make_section(late_ring, end_ring);
#        return end_ring

    # Recursive version of draw_branch
    def draw_branch(start, end, radius, prev_pos, subdivisions, start_ring):
        # Tuesday's bug:  I used end instead of prev_pos
        # Really, it was that simple!
        section_length = (start - end).length   # My fix actually had this part wrong - should be end
        continue_vector = section_length * (start - prev_pos).normalized()
        # Suggestion for small side branches:
        #   Scale radius of recursive calls based on how vertical (start - end) is

        middle_of_straight_branch = (start + end) / 2
        end_of_continue = start + continue_vector
        half_continue = start + continue_vector/2
        early_spot = (2 * start + middle_of_straight_branch + 2 * half_continue ) / 5
        # I did a little re-weighting on middle_spot and late_spot to avoid abrupt jogs
        middle_spot = (2 * start + end + end_of_continue) / 4
        late_spot = (middle_spot + 2 * end + middle_of_straight_branch) / 4
        # if random() * 5 > radius: # or something like that
        #     make_leaf(middle_spot, cross product of start - end vector and (0, 0, 1), 1)
        
        # We'll need the rings we calculated earlier to make the recursive call
        if subdivisions > 1:
            early_ring = draw_branch(start, early_spot, radius * 1.2, prev_pos, subdivisions - 1, start_ring)
            middle_ring = draw_branch(early_spot, middle_spot, radius * 1.1, start, subdivisions - 1, early_ring)
            late_ring = draw_branch(middle_spot, late_spot, radius * 1.05, early_spot, subdivisions - 1, middle_ring)
            end_ring = draw_branch(late_spot, end, radius, middle_spot, subdivisions - 1, late_ring)
        else:
            early_ring = make_end_ring(start, early_spot, radius * 1.2)
            middle_ring = make_end_ring(early_spot, middle_spot, radius * 1.1)
            late_ring = make_end_ring(middle_spot, late_spot, radius * 1.05)
            end_ring = make_end_ring(late_spot, end, radius)
            make_section(start_ring, early_ring);
            make_section(early_ring, middle_ring);
            make_section(middle_ring, late_ring)
            make_section(late_ring, end_ring);
        return end_ring

 #   def make_leaf(pos, direction, size):
        # How many points or what shape?
        # What texture?
            
 
    # params?
    def make_tree(pos, calls, prev_pos, prev_ring, stop_adjustment = 1):
        # Will we stop?
        if random() > stop_adjustment /calls:
            # Leaves?
            if pos.z > top_fruit.location.z:
                top_fruit.location = pos
            return
        branches = randint(1, 3)
        # were we a branch?
        if Vector((0, 0, 1)).dot((pos - prev_pos).normalized()) < 0.9:
           branch_ends = []
        else:
            branch_ends = [pos + Vector((0, 0, 2))]
        branch_ends.append(pos + 3 * Vector((0.5 - random(), 0.5 - random(), random())))
        branch_ends.append(pos + 3 * Vector((0.5 - random(), 0.5 - random(), random())))
        for new_endpoint in branch_ends:
            # new_endpoint = Vector((pos.x + 2 * (random() - 0.5), pos.y + 2 * (random() - 0.5), pos.z + 1 + random()))
            # Draw branch to new endpoint
            branch_size = Vector((0, 0, 1)).dot((new_endpoint - pos).normalized())
            new_end_ring = draw_branch(pos, new_endpoint, .3/calls, prev_pos, 2, prev_ring)
            if branch_size > 0.9:
                new_calls = calls + 1
            else:
                new_calls = calls + 4
            make_tree(new_endpoint, new_calls, pos, new_end_ring, stop_adjustment=stop_adjustment)
 
    start_ring = make_end_ring(Vector((0, 0, 0)), Vector((0, 0, 2)), 0.3, start_instead=True)
    #draw_branch(Vector((0, 0, 0)), Vector((2, 2, 2)), 0.2, Vector((0, 0, -2)), 2, start_ring)
    #draw_branch(Vector((0, 0, 0)), Vector((2, -2, 2)), 0.2, Vector((0, 0, -2)), 2, start_ring)
    #draw_branch(Vector((0, 0, 0)), Vector((-2, 2, 2)), 0.2, Vector((0, 0, -2)), 2, start_ring)
    #draw_branch(Vector((0, 0, 0)), Vector((-2, -2, 2)), 0.2, Vector((0, 0, -2)), 2, start_ring)
    make_tree(Vector((0, 0, 0)), 1, Vector((0, 0, -2)), start_ring, 4)
    #start_ring = make_end_ring(Vector((4, 0, 0)), Vector((4, 0, 2)), 0.3, start_instead=True)
    #draw_branch(Vector((4, 0, 0)), Vector((4, 0, 2)), 0.2, Vector((4, 0, -2)), 2, start_ring)
    #start_ring = make_end_ring(Vector((-4, 0, 0)), Vector((-4, 0, 2)), 0.03, start_instead=True)
    #draw_branch(Vector((-4, 0, 0)), Vector((-4, 0, 2)), 0.02, Vector((-4, 1, -2)), 2, start_ring)

    tree_mesh = bpy.data.meshes.new(name="TreeMesh")
    tree_mesh.from_pydata(vertices=verts, edges=[], faces=faces)
    add_texture_coordinates(tree_mesh)
    tree_mesh.validate(verbose=True) # I think this prints out errors if we have any
    
    tree_object = bpy.data.objects.new(name="Tree", object_data=tree_mesh)
    bpy.context.collection.objects.link(tree_object)

    material = bpy.data.materials.new(name="TreeBark")
    material.use_nodes = True
    tree_object.data.materials.append(material)
    bsdf_node= material.node_tree.nodes.get("Principled BSDF")
    texture_node = material.node_tree.nodes.new(type="ShaderNodeTexImage")
    texture_node.image = bpy.data.images.load("bark.jpg")
    material.node_tree.links.new(texture_node.outputs["Color"], bsdf_node.inputs["Base Color"])
    
    return tree_object, top_fruit


tree, tf = new_tree()

bpy.context.scene.frame_start = 1
bpy.context.scene.frame_end = 360
tree.rotation_euler = (0, 0, 0)
tree.location = (0, 0, 0)
tree.scale = (0, 0, 0)
tree.keyframe_insert(data_path="scale", frame=1)
tree.scale = (1, 1, 0)
tree.keyframe_insert(data_path="rotation_euler", frame=100)
tree.keyframe_insert(data_path="scale", frame=10)
tree.keyframe_insert(data_path="location", frame=100)
tree.scale = (1, 1, 1)
tree.keyframe_insert(data_path="scale", frame=100)
tree.location = (100, 0, 0)
tree.keyframe_insert(data_path="location", frame=360);
tree.rotation_euler = (0, 0, 8*pi)
tree.keyframe_insert(data_path="rotation_euler", frame=180)
tree.rotation_euler = (0, 0, 0)
tree.keyframe_insert(data_path="rotation_euler", frame=360)


tf.scale = (0, 0, 0)
tf.keyframe_insert(data_path="scale", frame=1)
tf.scale = (1, 1, 0)
tf.keyframe_insert(data_path="scale", frame=10)
tf.scale = (1, 1, 1)
tf.keyframe_insert(data_path="scale", frame=100)

tf.keyframe_insert(data_path="location", frame=100)
tf.location = (tf.location.x, tf.location.y, 0)
tf.keyframe_insert(data_path="location", frame=10)
tf.location = (0, 0, 0)
tf.keyframe_insert(data_path="location", frame=1)
tf.keyframe_insert(data_path="location", frame=360)

for fcurve in tree.animation_data.action.fcurves:
	for keyframe in fcurve.keyframe_points:
		keyframe.interpolation = "LINEAR"

for fcurve in tf.animation_data.action.fcurves:
	for keyframe in fcurve.keyframe_points:
		keyframe.interpolation = "LINEAR"