import bpy
import bmesh
import math
import random
from mathutils import Vector, Euler

random.seed(42)

def clear_scene():
    bpy.ops.object.select_all(action='SELECT')
    bpy.ops.object.delete()
    for mesh in bpy.data.meshes:
        bpy.data.meshes.remove(mesh)
    for mat in bpy.data.materials:
        bpy.data.materials.remove(mat)
    for curve in bpy.data.curves:
        bpy.data.curves.remove(curve)

def create_material(name, base_color, roughness=0.5, subsurface_weight=0.0):
    mat = bpy.data.materials.new(name=name)
    mat.use_nodes = True
    nodes = mat.node_tree.nodes
    principled = nodes.get("Principled BSDF")
    if principled:
        try:
            principled.inputs["Base Color"].default_value = tuple(base_color)
        except Exception:
            pass
        try:
            principled.inputs["Roughness"].default_value = float(roughness)
        except Exception:
            pass
        try:
            principled.inputs["Subsurface Weight"].default_value = float(subsurface_weight)
        except Exception:
            pass
        try:
            principled.inputs["Subsurface Scale"].default_value = 0.1
        except Exception:
            pass
    return mat

def assign_material(obj, mat):
    if obj and hasattr(obj, 'data') and obj.data and hasattr(obj.data, 'materials'):
        obj.data.materials.append(mat)

def create_stalk():
    stalk_height = 2.2
    segments = 24
    rings = 16
    
    bpy.ops.mesh.primitive_cylinder_add(
        vertices=segments,
        radius=0.04,
        depth=stalk_height,
        end_fill_type='NGON',
        location=(0, 0, stalk_height / 2)
    )
    stalk = bpy.context.active_object
    if not stalk:
        return None
    stalk.name = "Stalk"
    
    bm = bmesh.new()
    bm.from_mesh(stalk.data)
    
    bmesh.ops.subdivide_edges(bm, edges=[e for e in bm.edges if abs(e.verts[0].co.z - e.verts[1].co.z) > 0.01], cuts=rings)
    
    bm.verts.ensure_lookup_table()
    for v in bm.verts:
        z_norm = (v.co.z + stalk_height / 2) / stalk_height
        taper = 1.0 - 0.35 * z_norm
        dist = math.sqrt(v.co.x**2 + v.co.y**2)
        if dist > 0.001:
            v.co.x *= taper
            v.co.y *= taper
        node_bulge = 0.0
        for node_z in [0.25, 0.45, 0.65, 0.85, 1.05, 1.25, 1.45, 1.65, 1.85]:
            node_z_local = node_z - stalk_height / 2
            if abs(v.co.z - node_z_local) < 0.03:
                node_bulge = 0.003 * (1 - abs(v.co.z - node_z_local) / 0.03)
        if dist > 0.001:
            factor = 1 + node_bulge / dist
            v.co.x *= factor
            v.co.y *= factor
    
    bm.to_mesh(stalk.data)
    bm.free()
    
    for poly in stalk.data.polygons:
        poly.use_smooth = True
    
    stalk_mat = create_material("Stalk_Mat", (0.18, 0.32, 0.08, 1.0), roughness=0.75)
    assign_material(stalk, stalk_mat)
    
    return stalk

def create_single_leaf(length, width, droop, twist):
    verts = []
    faces = []
    res_length = 20
    res_width = 6
    
    for i in range(res_length + 1):
        t = i / res_length
        for j in range(res_width + 1):
            s = j / res_width - 0.5
            
            x = t * length
            width_factor = math.sin(t * math.pi) * (1 - 0.3 * t)
            y = s * width * width_factor * 2
            
            z = -droop * t * t * t
            z += 0.01 * math.sin(t * math.pi * 2) * abs(s) * 2
            
            midrib_depth = 0.008 * (1 - abs(s) * 3) * (1 - t * 0.5)
            if abs(s) < 0.33:
                z -= max(0, midrib_depth)
            
            wave = 0.005 * math.sin(t * 8) * math.sin(s * math.pi)
            z += wave
            
            angle = twist * t
            y_rot = y * math.cos(angle) - z * math.sin(angle)
            z_rot = y * math.sin(angle) + z * math.cos(angle)
            
            verts.append((x, y_rot, z_rot))
    
    for i in range(res_length):
        for j in range(res_width):
            v0 = i * (res_width + 1) + j
            v1 = v0 + 1
            v2 = v0 + res_width + 2
            v3 = v0 + res_width + 1
            faces.append((v0, v1, v2, v3))
    
    mesh = bpy.data.meshes.new("Leaf_Mesh")
    mesh.from_pydata(verts, [], faces)
    mesh.update()
    
    obj = bpy.data.objects.new("Leaf", mesh)
    bpy.context.collection.objects.link(obj)
    
    for poly in obj.data.polygons:
        poly.use_smooth = True
    
    return obj

def create_leaves():
    leaves = []
    leaf_mat = create_material("Leaf_Mat", (0.12, 0.38, 0.06, 1.0), roughness=0.55, subsurface_weight=0.15)
    
    leaf_data = [
        (0.25, 0.55, 0.07, 0.28, 0.15),
        (0.45, 0.65, 0.08, 0.22, 0.12),
        (0.65, 0.75, 0.09, 0.18, 0.10),
        (0.85, 0.80, 0.095, 0.15, 0.08),
        (1.05, 0.78, 0.09, 0.16, 0.10),
        (1.25, 0.72, 0.085, 0.20, 0.12),
        (1.45, 0.62, 0.075, 0.25, 0.15),
        (1.65, 0.50, 0.065, 0.30, 0.18),
        (1.85, 0.38, 0.05, 0.35, 0.20),
    ]
    
    for i, (z_pos, length, width, droop, twist) in enumerate(leaf_data):
        length *= (0.9 + random.random() * 0.2)
        droop *= (0.8 + random.random() * 0.4)
        
        leaf = create_single_leaf(length, width, droop, twist * (1 if random.random() > 0.5 else -1))
        leaf.name = f"Leaf_{i}"
        
        angle = (i * 137.5 + random.uniform(-15, 15)) * math.pi / 180
        
        t = z_pos / 2.2
        stalk_radius = 0.04 * (1.0 - 0.35 * t) + 0.005
        
        leaf.location = (
            stalk_radius * math.cos(angle),
            stalk_radius * math.sin(angle),
            z_pos
        )
        
        leaf.rotation_euler = Euler((
            math.radians(10 + random.uniform(-5, 5)),
            math.radians(-15 - random.uniform(0, 15)),
            angle + math.pi / 2
        ), 'XYZ')
        
        assign_material(leaf, leaf_mat)
        leaves.append(leaf)
    
    return leaves

def create_leaf_sheaths():
    sheaths = []
    sheath_mat = create_material("Sheath_Mat", (0.20, 0.35, 0.10, 1.0), roughness=0.65)
    
    for i, z_pos in enumerate([0.25, 0.45, 0.65, 0.85, 1.05, 1.25, 1.45, 1.65, 1.85]):
        t = z_pos / 2.2
        stalk_radius = 0.04 * (1.0 - 0.35 * t)
        
        bpy.ops.mesh.primitive_cylinder_add(
            vertices=12,
            radius=stalk_radius + 0.006,
            depth=0.12,
            location=(0, 0, z_pos + 0.03)
        )
        sheath = bpy.context.active_object
        if sheath:
            sheath.name = f"Sheath_{i}"
            
            bm = bmesh.new()
            bm.from_mesh(sheath.data)
            bm.verts.ensure_lookup_table()
            for v in bm.verts:
                z_local = v.co.z
                factor = 1.0 + 0.15 * (z_local / 0.06)
                v.co.x *= factor
                v.co.y *= factor
            bm.to_mesh(sheath.data)
            bm.free()
            
            assign_material(sheath, sheath_mat)
            sheaths.append(sheath)
    
    return sheaths

def create_tassel():
    tassel_parts = []
    tassel_z = 2.25
    tassel_mat = create_material("Tassel_Mat", (0.75, 0.68, 0.45, 1.0), roughness=0.8)
    
    bpy.ops.mesh.primitive_cone_add(
        vertices=8,
        radius1=0.02,
        radius2=0.008,
        depth=0.18,
        location=(0, 0, tassel_z + 0.09)
    )
    central = bpy.context.active_object
    if central:
        central.name = "Tassel_Central"
        assign_material(central, tassel_mat)
        tassel_parts.append(central)
    
    for tier in range(3):
        tier_z = tassel_z - 0.05 + tier * 0.04
        num_branches = 8 - tier * 2
        
        for i in range(num_branches):
            angle = (i / num_branches) * 2 * math.pi + tier * 0.3 + random.uniform(-0.1, 0.1)
            branch_length = 0.12 - tier * 0.02 + random.uniform(-0.02, 0.02)
            
            curve_data = bpy.data.curves.new(f"Tassel_Branch_{tier}_{i}", 'CURVE')
            curve_data.dimensions = '3D'
            curve_data.bevel_depth = 0.003 - tier * 0.0005
            curve_data.bevel_resolution = 2
            
            spline = curve_data.splines.new('BEZIER')
            spline.bezier_points.add(2)
            
            tilt = math.radians(40 + tier * 15 + random.uniform(-10, 10))
            
            for j, point in enumerate(spline.bezier_points):
                t = j / 2
                r = t * branch_length
                point.co = (
                    r * math.sin(tilt) * math.cos(angle),
                    r * math.sin(tilt) * math.sin(angle),
                    r * math.cos(tilt) + random.uniform(-0.01, 0.01)
                )
                point.handle_left_type = 'AUTO'
                point.handle_right_type = 'AUTO'
            
            branch = bpy.data.objects.new(f"Tassel_Branch_{tier}_{i}", curve_data)
            bpy.context.collection.objects.link(branch)
            branch.location = (0, 0, tier_z)
            assign_material(branch, tassel_mat)
            tassel_parts.append(branch)
            
            for k in range(3):
                spikelet_data = bpy.data.curves.new(f"Spikelet_{tier}_{i}_{k}", 'CURVE')
                spikelet_data.dimensions = '3D'
                spikelet_data.bevel_depth = 0.001
                
                sp_spline = spikelet_data.splines.new('BEZIER')
                sp_spline.bezier_points.add(1)
                
                t_pos = 0.4 + k * 0.25
                base_r = t_pos * branch_length
                base_x = base_r * math.sin(tilt) * math.cos(angle)
                base_y = base_r * math.sin(tilt) * math.sin(angle)
                base_z = base_r * math.cos(tilt)
                
                sp_angle = random.uniform(0, 2 * math.pi)
                sp_len = 0.02 + random.uniform(0, 0.01)
                
                sp_spline.bezier_points[0].co = (base_x, base_y, base_z)
                sp_spline.bezier_points[1].co = (
                    base_x + sp_len * math.cos(sp_angle),
                    base_y + sp_len * math.sin(sp_angle),
                    base_z + sp_len * 0.3
                )
                sp_spline.bezier_points[0].handle_left_type = 'AUTO'
                sp_spline.bezier_points[0].handle_right_type = 'AUTO'
                sp_spline.bezier_points[1].handle_left_type = 'AUTO'
                sp_spline.bezier_points[1].handle_right_type = 'AUTO'
                
                spikelet = bpy.data.objects.new(f"Spikelet_{tier}_{i}_{k}", spikelet_data)
                bpy.context.collection.objects.link(spikelet)
                spikelet.location = (0, 0, tier_z)
                assign_material(spikelet, tassel_mat)
                tassel_parts.append(spikelet)
    
    return tassel_parts

def create_ear():
    ear_parts = []
    ear_z = 0.95
    ear_angle = math.radians(25)
    ear_offset = 0.05
    
    cob_mat = create_material("Cob_Mat", (0.85, 0.72, 0.15, 1.0), roughness=0.45)
    
    bpy.ops.mesh.primitive_cylinder_add(
        vertices=24,
        radius=0.032,
        depth=0.22,
        location=(0, 0, 0)
    )
    cob = bpy.context.active_object
    if cob:
        cob.name = "Ear_Cob"
        
        bm = bmesh.new()
        bm.from_mesh(cob.data)
        
        bmesh.ops.subdivide_edges(bm, edges=[e for e in bm.edges], cuts=4)
        bm.verts.ensure_lookup_table()
        
        for v in bm.verts:
            z = v.co.z
            dist = math.sqrt(v.co.x**2 + v.co.y**2)
            if dist > 0.001:
                angle = math.atan2(v.co.y, v.co.x)
                
                end_taper = 1.0
                if abs(z) > 0.08:
                    end_taper = 1.0 - (abs(z) - 0.08) / 0.03
                    end_taper = max(0.3, end_taper)
                
                kernel_rows = 16
                kernel_bump = 0.003 * abs(math.sin(angle * kernel_rows))
                kernel_z_bump = 0.002 * abs(math.sin(z * 60))
                
                new_dist = (dist + kernel_bump + kernel_z_bump) * end_taper
                v.co.x = new_dist * math.cos(angle)
                v.co.y = new_dist * math.sin(angle)
        
        bm.to_mesh(cob.data)
        bm.free()
        
        cob.location = (ear_offset + 0.06, 0, ear_z)
        cob.rotation_euler = (0, ear_angle, 0)
        
        assign_material(cob, cob_mat)
        ear_parts.append(cob)
    
    husk_mat = create_material("Husk_Mat", (0.28, 0.42, 0.12, 1.0), roughness=0.6, subsurface_weight=0.08)
    
    for h in range(7):
        verts = []
        faces = []
        res_u = 12
        res_v = 16
        
        base_angle = h * (2 * math.pi / 7)
        coverage = 0.7 - h * 0.05
        
        for i in range(res_v + 1):
            v_t = i / res_v
            for j in range(res_u + 1):
                u_t = j / res_u
                
                z = (v_t - 0.5) * 0.28
                
                angle = base_angle + (u_t - 0.5) * coverage * math.pi
                
                radius = 0.038 + 0.004 * (1 - v_t) + h * 0.002
                radius *= (1 - 0.15 * abs(v_t - 0.5))
                
                if v_t > 0.85:
                    radius *= 1.0 - (v_t - 0.85) / 0.15 * 0.5
                
                x = radius * math.cos(angle)
                y = radius * math.sin(angle)
                
                verts.append((x, y, z))
        
        for i in range(res_v):
            for j in range(res_u):
                v0 = i * (res_u + 1) + j
                v1 = v0 + 1
                v2 = v0 + res_u + 2
                v3 = v0 + res_u + 1
                faces.append((v0, v1, v2, v3))
        
        mesh = bpy.data.meshes.new(f"Husk_Mesh_{h}")
        mesh.from_pydata(verts, [], faces)
        mesh.update()
        
        husk = bpy.data.objects.new(f"Husk_{h}", mesh)
        bpy.context.collection.objects.link(husk)
        
        husk.location = (ear_offset + 0.06, 0, ear_z)
        husk.rotation_euler = (0, ear_angle, 0)
        
        for poly in husk.data.polygons:
            poly.use_smooth = True
        
        assign_material(husk, husk_mat)
        ear_parts.append(husk)
    
    silk_mat = create_material("Silk_Mat", (0.82, 0.75, 0.55, 1.0), roughness=0.9)
    
    for s in range(30):
        curve_data = bpy.data.curves.new(f"Silk_{s}", 'CURVE')
        curve_data.dimensions = '3D'
        curve_data.bevel_depth = 0.0015
        curve_data.bevel_resolution = 1
        
        spline = curve_data.splines.new('BEZIER')
        spline.bezier_points.add(4)
        
        start_angle = random.uniform(0, 2 * math.pi)
        start_r = 0.02 + random.uniform(0, 0.01)
        
        silk_length = 0.08 + random.uniform(0, 0.06)
        
        for j, point in enumerate(spline.bezier_points):
            t = j / 4
            point.co = (
                start_r * math.cos(start_angle) + t * silk_length * random.uniform(0.5, 1.5),
                start_r * math.sin(start_angle) + t * silk_length * random.uniform(-0.3, 0.3),
                0.12 + t * silk_length * 0.8 + random.uniform(-0.01, 0.01)
            )
            point.handle_left_type = 'AUTO'
            point.handle_right_type = 'AUTO'
        
        silk = bpy.data.objects.new(f"Silk_{s}", curve_data)
        bpy.context.collection.objects.link(silk)
        silk.location = (ear_offset + 0.06, 0, ear_z)
        silk.rotation_euler = (0, ear_angle, 0)
        assign_material(silk, silk_mat)
        ear_parts.append(silk)
    
    peduncle_mat = create_material("Peduncle_Mat", (0.22, 0.35, 0.10, 1.0), roughness=0.7)
    
    bpy.ops.mesh.primitive_cylinder_add(
        vertices=8,
        radius=0.015,
        depth=0.08,
        location=(ear_offset / 2, 0, ear_z - 0.02)
    )
    peduncle = bpy.context.active_object
    if peduncle:
        peduncle.name = "Ear_Peduncle"
        peduncle.rotation_euler = (0, ear_angle, 0)
        assign_material(peduncle, peduncle_mat)
        ear_parts.append(peduncle)
    
    return ear_parts

def create_roots():
    root_parts = []
    root_mat = create_material("Root_Mat", (0.38, 0.28, 0.18, 1.0), roughness=0.85)
    
    for r in range(16):
        curve_data = bpy.data.curves.new(f"Root_{r}", 'CURVE')
        curve_data.dimensions = '3D'
        curve_data.bevel_depth = 0.006 - r * 0.0002
        curve_data.bevel_resolution = 2
        
        spline = curve_data.splines.new('BEZIER')
        spline.bezier_points.add(5)
        
        angle = r * (2 * math.pi / 16) + random.uniform(-0.2, 0.2)
        root_length = 0.25 + random.uniform(0, 0.15)
        spread = 0.12 + random.uniform(0, 0.08)
        
        for j, point in enumerate(spline.bezier_points):
            t = j / 5
            point.co = (
                t * spread * math.cos(angle) + random.uniform(-0.01, 0.01),
                t * spread * math.sin(angle) + random.uniform(-0.01, 0.01),
                -t * root_length - t * t * 0.05
            )
            point.handle_left_type = 'AUTO'
            point.handle_right_type = 'AUTO'
        
        root = bpy.data.objects.new(f"Root_{r}", curve_data)
        bpy.context.collection.objects.link(root)
        assign_material(root, root_mat)
        root_parts.append(root)
        
        for br in range(2):
            branch_data = bpy.data.curves.new(f"Root_{r}_Branch_{br}", 'CURVE')
            branch_data.dimensions = '3D'
            branch_data.bevel_depth = 0.003
            branch_data.bevel_resolution = 1
            
            br_spline = branch_data.splines.new('BEZIER')
            br_spline.bezier_points.add(3)
            
            start_t = 0.3 + br * 0.3
            start_x = start_t * spread * math.cos(angle)
            start_y = start_t * spread * math.sin(angle)
            start_z = -start_t * root_length
            
            br_angle = angle + random.uniform(-0.8, 0.8)
            br_length = 0.08 + random.uniform(0, 0.05)
            
            for j, point in enumerate(br_spline.bezier_points):
                t = j / 3
                point.co = (
                    start_x + t * br_length * math.cos(br_angle),
                    start_y + t * br_length * math.sin(br_angle),
                    start_z - t * br_length * 0.6
                )
                point.handle_left_type = 'AUTO'
                point.handle_right_type = 'AUTO'
            
            branch = bpy.data.objects.new(f"Root_{r}_Branch_{br}", branch_data)
            bpy.context.collection.objects.link(branch)
            assign_material(branch, root_mat)
            root_parts.append(branch)
    
    return root_parts

def create_brace_roots():
    brace_parts = []
    brace_mat = create_material("Brace_Root_Mat", (0.35, 0.26, 0.16, 1.0), roughness=0.8)
    
    for tier in range(2):
        tier_z = 0.08 + tier * 0.08
        num_roots = 6 - tier
        
        for r in range(num_roots):
            curve_data = bpy.data.curves.new(f"Brace_Root_{tier}_{r}", 'CURVE')
            curve_data.dimensions = '3D'
            curve_data.bevel_depth = 0.008 - tier * 0.002
            curve_data.bevel_resolution = 3
            
            spline = curve_data.splines.new('BEZIER')
            spline.bezier_points.add(4)
            
            angle = r * (2 * math.pi / num_roots) + tier * 0.3 + random.uniform(-0.15, 0.15)
            
            for j, point in enumerate(spline.bezier_points):
                t = j / 4
                radius = t * 0.10
                z = tier_z * (1 - t * t)
                point.co = (
                    radius * math.cos(angle),
                    radius * math.sin(angle),
                    z
                )
                point.handle_left_type = 'AUTO'
                point.handle_right_type = 'AUTO'
            
            brace = bpy.data.objects.new(f"Brace_Root_{tier}_{r}", curve_data)
            bpy.context.collection.objects.link(brace)
            assign_material(brace, brace_mat)
            brace_parts.append(brace)
    
    return brace_parts

def create_internode_details():
    details = []
    
    for i, z in enumerate([0.35, 0.55, 0.75, 0.95, 1.15, 1.35, 1.55, 1.75]):
        t = z / 2.2
        radius = 0.04 * (1.0 - 0.35 * t)
        
        groove_mat = create_material(f"Groove_Mat_{i}", (0.14, 0.28, 0.06, 1.0), roughness=0.7)
        
        bpy.ops.mesh.primitive_torus_add(
            major_radius=radius + 0.002,
            minor_radius=0.002,
            major_segments=24,
            minor_segments=6,
            location=(0, 0, z)
        )
        groove = bpy.context.active_object
        if groove:
            groove.name = f"Internode_Ring_{i}"
            assign_material(groove, groove_mat)
            details.append(groove)
    
    return details

def main():
    clear_scene()
    
    all_parts = []
    
    stalk = create_stalk()
    if stalk:
        all_parts.append(stalk)
    
    sheaths = create_leaf_sheaths()
    all_parts.extend(sheaths)
    
    details = create_internode_details()
    all_parts.extend(details)
    
    leaves = create_leaves()
    all_parts.extend(leaves)
    
    tassel = create_tassel()
    all_parts.extend(tassel)
    
    ear = create_ear()
    all_parts.extend(ear)
    
    roots = create_roots()
    all_parts.extend(roots)
    
    brace_roots = create_brace_roots()
    all_parts.extend(brace_roots)
    
    bpy.ops.object.empty_add(type='PLAIN_AXES', location=(0, 0, 0))
    parent = bpy.context.active_object
    if parent:
        parent.name = "Zea_mays"
        
        for obj in all_parts:
            if obj:
                try:
                    obj.parent = parent
                except Exception:
                    pass
        
        bpy.ops.object.select_all(action='DESELECT')
        parent.select_set(True)
        bpy.context.view_layer.objects.active = parent

if __name__ == "__main__":
    main()