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

random.seed(20260925)

# ----------------------------------------------------------------------
# Clean start (safe): remove everything currently in the file.
# ----------------------------------------------------------------------
for ob in list(bpy.data.objects):
    try:
        bpy.data.objects.remove(ob, do_unlink=True)
    except Exception:
        pass

ALL_PARTS = []

# ----------------------------------------------------------------------
# Helpers
# ----------------------------------------------------------------------
def obj_from_bmesh(name, bm):
    me = bpy.data.meshes.new(name)
    try:
        bm.normal_update()
    except Exception:
        pass
    bm.to_mesh(me)
    bm.free()
    ob = bpy.data.objects.new(name, me)
    bpy.context.collection.objects.link(ob)
    return ob


def shade_smooth(ob):
    try:
        for p in ob.data.polygons:
            p.use_smooth = True
    except Exception:
        pass


def make_material(name, color, roughness=0.6, subsurface=0.0):
    mat = bpy.data.materials.new(name)
    mat.use_nodes = True
    try:
        nt = mat.node_tree
        bsdf = None
        for n in nt.nodes:
            if n.type == 'BSDF_PRINCIPLED':
                bsdf = n
                break
        if bsdf is not None:
            def set_in(key, val):
                try:
                    if key in bsdf.inputs:
                        bsdf.inputs[key].default_value = val
                except Exception:
                    pass
            set_in("Base Color", (color[0], color[1], color[2], 1.0))
            set_in("Roughness", roughness)
            # 4.x renamed socket -> guard both spellings defensively
            if subsurface > 0.0:
                set_in("Subsurface Weight", subsurface)
                set_in("Subsurface", subsurface)
                set_in("Subsurface Radius", (0.1, 0.2, 0.05))
    except Exception:
        pass
    return mat


def assign_mat(ob, mat):
    try:
        ob.data.materials.append(mat)
    except Exception:
        pass


def cyl(name, r1, r2, depth, segments=14):
    bm = bmesh.new()
    bmesh.ops.create_cone(bm, cap_ends=True, cap_tris=False,
                          segments=segments, radius1=r1, radius2=r2, depth=depth)
    ob = obj_from_bmesh(name, bm)
    shade_smooth(ob)
    return ob


# Materials -------------------------------------------------------------
mat_stem   = make_material("Maize_Stem",   (0.24, 0.42, 0.11), 0.55, 0.15)
mat_leaf   = make_material("Maize_Leaf",   (0.15, 0.40, 0.08), 0.50, 0.25)
mat_husk   = make_material("Maize_Husk",   (0.52, 0.56, 0.24), 0.60, 0.15)
mat_kernel = make_material("Maize_Kernel", (0.88, 0.70, 0.10), 0.35)
mat_cob    = make_material("Maize_Cob",    (0.92, 0.86, 0.62), 0.60)
mat_silk   = make_material("Maize_Silk",   (0.80, 0.72, 0.42), 0.70)
mat_root   = make_material("Maize_Root",   (0.62, 0.50, 0.34), 0.80)
mat_tassel = make_material("Maize_Tassel", (0.72, 0.66, 0.34), 0.55)


# ----------------------------------------------------------------------
# Leaf builder  (local: base at origin, blade extends +X)
# ----------------------------------------------------------------------
def make_leaf(name, length, max_width, droop=0.9, arch=0.16):
    bm = bmesh.new()
    nseg = 26
    ncross = 7
    grid = [[None] * ncross for _ in range(nseg + 1)]
    for i in range(nseg + 1):
        s = i / nseg
        x = length * s
        # centerline: gentle rise then droop toward tip
        z = math.sin(min(s, 1.0) * math.pi) * length * arch - droop * length * (s ** 2) * 0.35
        # width profile: narrow sheath, broad middle, tapering point
        w = max_width * (math.sin(min(s * 1.02, 1.0) * math.pi) ** 0.55)
        if s > 0.55:
            w *= (1.0 - (s - 0.55) * 0.55)
        w = max(w, max_width * 0.02)
        for j in range(ncross):
            t = j / (ncross - 1) - 0.5
            y = t * w
            rib = (1.0 - abs(t) * 2.0)          # 1 at midrib, 0 at edge
            zz = z + rib * max_width * 0.10 - abs(t) * w * 0.12
            grid[i][j] = bm.verts.new((x, y, zz))
    for i in range(nseg):
        for j in range(ncross - 1):
            try:
                bm.faces.new((grid[i][j], grid[i][j + 1],
                              grid[i + 1][j + 1], grid[i + 1][j]))
            except Exception:
                pass
    ob = obj_from_bmesh(name, bm)
    shade_smooth(ob)
    return ob


# ----------------------------------------------------------------------
# Ear builder  (cob points +X, base at origin) -> returns empty group
# ----------------------------------------------------------------------
def make_ear(name):
    parts = []
    cob_len = 0.22
    cob_r = 0.028

    # Cob core (along +X)
    cob = cyl(name + "_Cob", cob_r * 0.35, cob_r, cob_len, segments=16)
    cob.rotation_euler = (0.0, math.radians(90.0), 0.0)
    cob.location = (cob_len * 0.5, 0.0, 0.0)
    assign_mat(cob, mat_cob)
    parts.append(cob)

    # Kernels (all packed into one mesh)
    bm = bmesh.new()
    rows_around = 16
    rows_along = 22
    for a in range(rows_along):
        xf = 0.02 + (cob_len - 0.04) * (a / (rows_along - 1))
        rr = cob_r * (0.55 + 0.45 * math.sin(math.pi * (a + 0.5) / rows_along))
        offset = (a % 2) * (math.pi / rows_around)
        for k in range(rows_around):
            ang = 2.0 * math.pi * k / rows_around + offset
            y = (rr + 0.004) * math.cos(ang)
            z = (rr + 0.004) * math.sin(ang)
            m = Matrix.Translation((xf, y, z))
            bmesh.ops.create_icosphere(bm, subdivisions=1, radius=0.007, matrix=m)
    kern = obj_from_bmesh(name + "_Kernels", bm)
    shade_smooth(kern)
    assign_mat(kern, mat_kernel)
    parts.append(kern)

    # Husk leaves wrapping the cob
    for h in range(6):
        hk = make_leaf(name + "_Husk_%d" % h, cob_len * 1.35, cob_r * 1.7,
                       droop=0.4, arch=0.05)
        assign_mat(hk, mat_husk)
        roll = 2.0 * math.pi * h / 6 + random.uniform(-0.15, 0.15)
        # wrap slightly outward around the cob axis (X)
        hk.rotation_euler = (roll, random.uniform(-0.05, 0.05), 0.0)
        hk.location = (0.0,
                       (cob_r * 1.05) * math.cos(roll),
                       (cob_r * 1.05) * math.sin(roll))
        parts.append(hk)

    # Silk strands at the tip (one mesh)
    bm = bmesh.new()
    for s in range(16):
        ang = 2.0 * math.pi * s / 16
        tilt = random.uniform(0.3, 0.7)
        rot = (Matrix.Rotation(ang, 4, 'X')
               @ Matrix.Rotation(math.radians(90.0), 4, 'Y')
               @ Matrix.Rotation(tilt, 4, 'X'))
        pos = Matrix.Translation((cob_len + 0.05, 0.0, 0.0))
        m = pos @ rot
        bmesh.ops.create_cone(bm, cap_ends=True, cap_tris=False, segments=5,
                              radius1=0.0016, radius2=0.0004, depth=0.11,
                              matrix=m)
    silk = obj_from_bmesh(name + "_Silk", bm)
    shade_smooth(silk)
    assign_mat(silk, mat_silk)
    parts.append(silk)

    # Group the ear parts under an empty
    emp = bpy.data.objects.new(name + "_Group", None)
    bpy.context.collection.objects.link(emp)
    for p in parts:
        p.parent = emp
    ALL_PARTS.extend(parts)
    ALL_PARTS.append(emp)
    return emp


# ----------------------------------------------------------------------
# Tassel builder (male inflorescence at top)  local base at origin, +Z up
# ----------------------------------------------------------------------
def make_tassel(name):
    parts = []
    rachis_h = 0.35
    rachis = cyl(name + "_Rachis", 0.010, 0.003, rachis_h, segments=10)
    rachis.location = (0.0, 0.0, rachis_h * 0.5)
    assign_mat(rachis, mat_tassel)
    parts.append(rachis)

    nbranch = 8
    spikelet_bm = bmesh.new()
    for b in range(nbranch):
        blen = random.uniform(0.12, 0.22)
        az = 2.0 * math.pi * b / nbranch + random.uniform(-0.2, 0.2)
        z0 = rachis_h * (0.15 + 0.6 * (b / nbranch))
        tilt = random.uniform(0.5, 0.9)
        br = cyl("%s_Branch_%d" % (name, b), 0.004, 0.0015, blen, segments=6)
        rot = Euler((0.0, tilt, az), 'XYZ')
        br.rotation_euler = rot
        # place branch base on the rachis, extend outward/up
        dirv = rot.to_matrix() @ Vector((0.0, 0.0, blen * 0.5))
        br.location = (dirv.x, dirv.y, z0 + dirv.z)
        assign_mat(br, mat_tassel)
        parts.append(br)
        # spikelets scattered along each branch
        M = rot.to_matrix().to_4x4()
        for sp in range(6):
            f = sp / 6.0
            local = Vector((0.0, 0.0, (f - 0.5) * blen))
            wpos = (Matrix.Translation((0.0, 0.0, z0)) @ M
                    @ Matrix.Translation((0.0, 0.0, blen * 0.5))).to_translation()
            base = M @ local
            m = Matrix.Translation((base.x, base.y, z0 + blen * 0.5 + local.z))
            bmesh.ops.create_icosphere(spikelet_bm, subdivisions=1,
                                       radius=0.005, matrix=m)
    spikes = obj_from_bmesh(name + "_Spikelets", spikelet_bm)
    shade_smooth(spikes)
    assign_mat(spikes, mat_tassel)
    parts.append(spikes)

    emp = bpy.data.objects.new(name + "_Group", None)
    bpy.context.collection.objects.link(emp)
    for p in parts:
        p.parent = emp
    ALL_PARTS.extend(parts)
    ALL_PARTS.append(emp)
    return emp


# ======================================================================
# BUILD THE WHOLE PLANT
# ======================================================================
STALK_H = 2.35
r_base = 0.055
r_top = 0.016

# --- Main stalk (tapering, with node bulges) ---
n_nodes = 9
seg_bm = bmesh.new()
rings = 40
for i in range(rings + 1):
    t = i / rings
    z = STALK_H * t
    r = r_base + (r_top - r_base) * t
    # node bulges
    node_phase = (t * n_nodes) % 1.0
    bulge = 1.0 + 0.22 * math.exp(-((node_phase - 0.0) ** 2) / 0.004)
    bulge = max(bulge, 1.0 + 0.22 * math.exp(-((node_phase - 1.0) ** 2) / 0.004))
    rr = r * bulge
    ring = []
    seg = 16
    for j in range(seg):
        a = 2.0 * math.pi * j / seg
        v = seg_bm.verts.new((rr * math.cos(a), rr * math.sin(a), z))
        ring.append(v)
    if i == 0:
        first_ring = ring
    if i > 0:
        for j in range(seg):
            j2 = (j + 1) % seg
            try:
                seg_bm.faces.new((prev_ring[j], prev_ring[j2], ring[j2], ring[j]))
            except Exception:
                pass
    prev_ring = ring
    last_ring = ring
# caps
try:
    seg_bm.faces.new(list(reversed(first_ring)))
except Exception:
    pass
try:
    seg_bm.faces.new(last_ring)
except Exception:
    pass
stalk = obj_from_bmesh("Maize_Stalk", seg_bm)
shade_smooth(stalk)
assign_mat(stalk, mat_stem)
ALL_PARTS.append(stalk)

# --- Leaves along the stalk (alternate / distichous) ---
leaf_nodes = [0.10, 0.22, 0.35, 0.48, 0.60, 0.71, 0.81, 0.90]
for idx, frac in enumerate(leaf_nodes):
    z = STALK_H * frac
    r = r_base + (r_top - r_base) * frac
    # size peaks in mid-canopy
    peak = 1.0 - abs(frac - 0.5) * 1.1
    length = 0.65 + 0.75 * max(peak, 0.1)
    width = 0.06 + 0.05 * max(peak, 0.1)
    yaw = math.pi * (idx % 2) + random.uniform(-0.12, 0.12)
    lf = make_leaf("Maize_Leaf_%d" % idx, length, width,
                   droop=0.9 + random.uniform(-0.1, 0.2), arch=0.17)
    assign_mat(lf, mat_leaf)
    pitch = -math.radians(38.0 + random.uniform(-6, 6))
    lf.rotation_euler = Euler((0.0, pitch, yaw), 'XYZ')
    lf.location = (r * 0.7 * math.cos(yaw), r * 0.7 * math.sin(yaw), z)
    ALL_PARTS.append(lf)

# --- One ear at a mid node ---
ear_frac = 0.44
ez = STALK_H * ear_frac
er = r_base + (r_top - r_base) * ear_frac
ear_yaw = math.radians(35.0)
ear = make_ear("Maize_Ear")
ear.rotation_euler = Euler((0.0, -math.radians(45.0), ear_yaw), 'XYZ')
ear.location = (er * math.cos(ear_yaw), er * math.sin(ear_yaw), ez)

# --- Tassel at the very top ---
tassel = make_tassel("Maize_Tassel")
tassel.location = (0.0, 0.0, STALK_H)

# --- Roots ---
# Fibrous seminal / crown roots at base
for i in range(14):
    az = 2.0 * math.pi * i / 14 + random.uniform(-0.15, 0.15)
    rlen = random.uniform(0.18, 0.34)
    root = cyl("Maize_Root_%d" % i, 0.010, 0.002, rlen, segments=6)
    spread = random.uniform(0.7, 1.15)  # radians from vertical
    root.rotation_euler = Euler((0.0, math.pi - spread, az), 'XYZ')
    dv = root.rotation_euler.to_matrix() @ Vector((0.0, 0.0, rlen * 0.5))
    root.location = (dv.x, dv.y, dv.z)
    assign_mat(root, mat_root)
    ALL_PARTS.append(root)

# Prop / brace roots arching from the lowest above-ground node
brace_z = STALK_H * 0.05
for i in range(6):
    az = 2.0 * math.pi * i / 6 + 0.3
    blen = random.uniform(0.14, 0.2)
    br = cyl("Maize_BraceRoot_%d" % i, 0.008, 0.003, blen, segments=6)
    spread = 1.0
    br.rotation_euler = Euler((0.0, math.pi - spread, az), 'XYZ')
    dv = br.rotation_euler.to_matrix() @ Vector((0.0, 0.0, blen * 0.5))
    br.location = (r_base * math.cos(az) + dv.x,
                   r_base * math.sin(az) + dv.y,
                   brace_z + dv.z)
    assign_mat(br, mat_root)
    ALL_PARTS.append(br)

# ----------------------------------------------------------------------
# Parent everything under a single root empty = one whole specimen
# ----------------------------------------------------------------------
root_empty = bpy.data.objects.new("Maize_Zea_mays", None)
bpy.context.collection.objects.link(root_empty)
root_empty.location = (0.0, 0.0, 0.0)

ear.parent = root_empty
tassel.parent = root_empty
for p in ALL_PARTS:
    if p is ear or p is tassel:
        continue
    if p.parent is None:
        p.parent = root_empty

try:
    bpy.context.view_layer.update()
except Exception:
    pass