import bpy
import bmesh
from mathutils import Vector
from math import *
from test_this import *

print_message()

def setup_rotation_animation(obj):
    bpy.context.scene.frame_start = 1
    bpy.context.scene.frame_end = 360
    obj.rotation_euler = (0, 0, 0)
    obj.keyframe_insert(data_path="rotation_euler", frame=1)
    obj.rotation_euler = (0, 0, 2*pi)
    obj.keyframe_insert(data_path="rotation_euler", frame=360)
    for fcurve in obj.animation_data.action.fcurves:
    	for keyframe in fcurve.keyframe_points:
    		keyframe.interpolation = "LINEAR"

def make_vertex_color_material():
    mat = bpy.data.materials.new(name="VertexColorMaterial")
    mat.use_nodes = True
    bsdf = mat.node_tree.nodes.get("Principled BSDF")
    color_node = mat.node_tree.nodes.new(type='ShaderNodeVertexColor')
    color_node.layer_name = "Col"
    mat.node_tree.links.new(color_node.outputs['Color'], bsdf.inputs['Base Color'])
    return mat

def delete_all_objects():
    bpy.ops.object.select_all(action='DESELECT')
    bpy.ops.object.select_by_type(type='MESH')
    bpy.ops.object.delete()

def set_material_view_and_zoom(zoom_level=1.0):
    for area in bpy.context.screen.areas:
    	if area.type == "VIEW_3D":
    		for space in area.spaces:
    			if space.type == "VIEW_3D":
    				space.shading.type = "MATERIAL"
    				space.region_3d.view_distance *= zoom_level

# This part is for the snowflake in particular

def get_edge_key(v1, v2):
	def vec_key(v):
		return (round(v.x, 6), round(v.y, 6), round(v.z, 6))
	return tuple((vec_key(v1), vec_key(v2)))

def point_on_edge(v1, v2, t):
	return v1.lerp(v2, t).copy()  # Ensure a new Vector

def generate_koch_triangle(v1, v2, v3, edge_points):
	"""
	Given triangle vertices and global edge point cache, return smaller triangle verts.
	"""
	# Edge keys to ensure shared subdivisions
	e1 = get_edge_key(v1, v2)
	e2 = get_edge_key(v2, v3)
	e3 = get_edge_key(v3, v1)

	# Get or create subdivision points
	def get_third_points(key, a, b):
		if key not in edge_points:
			p1 = point_on_edge(a, b, 1/3)
			p2 = point_on_edge(a, b, 2/3)
			edge_points[key] = (p1, p2)
		return edge_points[key]

	a1, a2 = get_third_points(e1, v1, v2)
	b1, b2 = get_third_points(e2, v2, v3)
	c1, c2 = get_third_points(e3, v3, v1)

	# Center "bump" point for Koch-style pattern
	center = (v1 + v2 + v3) / 3
	normal = (v2 - v1).cross(v3 - v1).normalized()
	bump = (center + normal * ((v2 - v1).length / 1)).copy()

	# Return list of triangle vertex triplets
	return [
		(v1.copy(), a1, c2),
		(a1, a2, bump),
		(b1, b2, bump),
		(c1, c2, bump),
		(a2, v2.copy(), b1),
		(c1, b2, v3.copy()),
		(c2, a1, bump),
		(a2, b1, bump),
		(b2, c1, bump)
	]

level = 3

delete_all_objects()

# Create Icosphere
bpy.ops.mesh.primitive_ico_sphere_add(subdivisions=2, radius=1)
obj = bpy.context.active_object
obj.name = "bob"
mesh = obj.data

# Switch to Edit mode to access BMesh
bpy.ops.object.mode_set(mode='EDIT')
bm = bmesh.from_edit_mesh(mesh)
bm.faces.ensure_lookup_table()

# Prepare edge cache and vertex cache
edge_points = {}
vert_cache = {}

# Store all triangles with coordinates
triangles = []
for face in bm.faces:
	if len(face.verts) == 3:
		triangles.append(tuple(v.co.copy() for v in face.verts))

# Clear original geometry
bmesh.ops.delete(bm, geom=bm.faces[:], context='FACES_ONLY')

# Helper to avoid duplicate vertices
def get_vert(co):
	key = tuple(round(c, 6) for c in co)
	if key not in vert_cache:
		vert_cache[key] = bm.verts.new(co)
	return vert_cache[key]

# Subdivide each triangle
for v1, v2, v3 in triangles:
	tris = generate_koch_triangle(v1, v2, v3, edge_points)
	for i in range(level):
		new_tris = []
		for tri in tris:
			new_tris.extend(generate_koch_triangle(tri[0], tri[1], tri[2], edge_points))
		tris = new_tris

	for tri in tris:
		verts = [get_vert(v) for v in tri]
		try:
			bm.faces.new(verts)
		except ValueError:
			pass  # Face already exists


obj.data.materials.append(make_vertex_color_material())

# Done with the snowflake

# Coloring (for the snowflake)
color_layer = bm.loops.layers.color.new("Col")
face_color = (1, 1, 1, 1)
for i, face in enumerate(bm.faces):
	for loop in face.loops:
		d_center = loop.vert.co.length
		loop[color_layer] = (0.5 + sin(pi/2 + 20.0 * d_center)/2, 0.5 + sin(16.0 * d_center)/2, 0.5 + sin(12 * d_center)/2, 1.0)

bmesh.update_edit_mesh(obj.data)
bpy.ops.object.mode_set(mode='OBJECT')

setup_rotation_animation(obj)
set_material_view_and_zoom(0.2)
bpy.ops.object.select_all(action='DESELECT')
bpy.ops.screen.animation_play()