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

random.seed(7)
TWO_PI = math.pi * 2.0

# ---------------------------------------------------------------------------
# Clean start: remove any pre-existing mesh objects (default cube etc.)
# (We only clear at the START, never at the end.)
# ---------------------------------------------------------------------------
bpy.ops.object.select_all(action='DESELECT')
for ob in list(bpy.data.objects):
    if ob.type == 'MESH':
        bpy.data.objects.remove(ob, do_unlink=True)

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def make_material(name, color, rough=0.75, subsurf=0.0, sub_color=None):
    mat = bpy.data.materials.new(name)
    mat.use_nodes = True
    nt = mat.node_tree
    bsdf = nt.nodes.get("Principled BSDF")
    if bsdf is not None:
        try:
            bsdf.inputs["Base Color"].default_value = (color[0], color[1], color[2], 1.0)
        except Exception:
            pass
        try:
            bsdf.inputs["Roughness"].default_value = float(rough)
        except Exception:
            pass
        if subsurf > 0.0:
            for sock in ("Subsurface Weight", "Subsurface"):
                if sock in bsdf.inputs:
                    try:
                        bsdf.inputs[sock].default_value = float(subsurf)
                    except Exception:
                        pass
                    break
            if sub_color is not None and "Subsurface Radius" in bsdf.inputs:
                try:
                    bsdf.inputs["Subsurface Radius"].default_value = sub_color
                except Exception:
                    pass
    return mat


def new_mesh_object(name, verts, faces, mat=None, smooth=True):
    mesh = bpy.data.meshes.new(name)
    mesh.from_pydata(verts, [], faces)
    mesh.update()
    if smooth:
        for p in mesh.polygons:
            p.use_smooth = True
    obj = bpy.data.objects.new(name, mesh)
    bpy.context.collection.objects.link(obj)
    if mat is not None:
        obj.data.materials.append(mat)
    return obj


# ---------------------------------------------------------------------------
# Materials
# ---------------------------------------------------------------------------
bark_mat  = make_material("Welwitschia_Bark",  (0.20, 0.13, 0.08), rough=0.9)
leaf_mat  = make_material("Welwitschia_Leaf",  (0.16, 0.28, 0.10), rough=0.6,
                          subsurf=0.15, sub_color=(0.4, 0.6, 0.2))
stalk_mat = make_material("Welwitschia_Stalk", (0.35, 0.22, 0.14), rough=0.8)
cone_mat  = make_material("Welwitschia_Cone",  (0.62, 0.22, 0.14), rough=0.55)

created_objects = []

# ---------------------------------------------------------------------------
# TRUNK  + WOODY CONCAVE CROWN  + TAPROOT  (single mesh)
# ---------------------------------------------------------------------------
R = 1.0
SEG = 48

verts = []
faces = []


def make_ring(rscale, z, hump=0.0, xk=1.12, yk=0.9):
    ring = []
    for j in range(SEG):
        th = TWO_PI * j / SEG
        x = rscale * xk * R * math.cos(th)
        y = rscale * yk * R * math.sin(th)
        zz = z + hump * abs(math.sin(th))
        ring.append(len(verts))
        verts.append((x, y, zz))
    return ring


def bridge(a, b):
    for j in range(SEG):
        j2 = (j + 1) % SEG
        faces.append((a[j], a[j2], b[j2], b[j]))


def fan(ring, apex):
    for j in range(SEG):
        j2 = (j + 1) % SEG
        faces.append((ring[j], ring[j2], apex))


# outer body (top rim -> down the woody obconical stem)
rim = make_ring(1.00, 0.52, hump=0.18)      # humped rim, two lips at +/-Y
b1  = make_ring(0.97, 0.38)
b2  = make_ring(0.85, 0.15)
b3  = make_ring(0.62, -0.10)
b4  = make_ring(0.35, -0.35)                 # base of trunk

# taproot rings (rounder, tapering deep down)
t1 = make_ring(0.30, -0.70, xk=1.0, yk=1.0)
t2 = make_ring(0.22, -1.30, xk=1.0, yk=1.0)
t3 = make_ring(0.15, -2.00, xk=1.0, yk=1.0)
t4 = make_ring(0.08, -2.80, xk=1.0, yk=1.0)

root_tip = len(verts)
verts.append((0.0, 0.0, -3.45))

# concave crown interior
inner  = make_ring(0.55, 0.28)
crown_center = len(verts)
verts.append((0.0, 0.0, 0.20))

# connect body
bridge(rim, b1)
bridge(b1, b2)
bridge(b2, b3)
bridge(b3, b4)
# connect root
bridge(b4, t1)
bridge(t1, t2)
bridge(t2, t3)
bridge(t3, t4)
fan(t4, root_tip)
# crown depression
bridge(rim, inner)
fan(inner, crown_center)

trunk_obj = new_mesh_object("Welwitschia_Trunk", verts, faces, bark_mat, smooth=True)
try:
    sub = trunk_obj.modifiers.new("Subsurf", 'SUBSURF')
    sub.levels = 1
    sub.render_levels = 2
except Exception:
    pass
created_objects.append(trunk_obj)

# ---------------------------------------------------------------------------
# THE TWO PERMANENT LEAVES (each split into ribbon-like strips, trailing on ground)
# ---------------------------------------------------------------------------
def build_leaf(sign):
    lverts = []
    lfaces = []
    base_y = sign * 0.90 * R * 0.95
    base_z = 0.50
    n_rib = 9
    nseg = 34
    half_base = 0.75
    for ri in range(n_rib):
        s = (ri / (n_rib - 1)) * 2.0 - 1.0          # -1 .. 1 across the strap
        x0 = s * half_base
        x_tip = s * 1.70
        length = 3.8 * (1.0 - 0.12 * abs(s))
        w0 = 0.16 * (1.0 - 0.25 * abs(s))
        phase = ri * 1.3 + (0.0 if sign > 0 else 0.7)
        left = []
        right = []
        for k in range(nseg + 1):
            t = k / nseg
            yy = base_y + sign * (0.10 + t * length)
            # arch up out of crown then flop down onto the ground
            if t < 0.30:
                zc = base_z * (1.0 - t / 0.30)
            else:
                zc = 0.0
            zc += 0.05 * math.sin(t * 9.0 + phase) * min(1.0, t / 0.30)  # tattered ground ripple
            if zc < 0.0:
                zc = 0.0
            xc = x0 + (x_tip - x0) * t + 0.12 * math.sin(t * 3.0 + phase) * t
            w = w0 * (1.0 - 0.55 * t)
            li = len(lverts); lverts.append((xc - w, yy, zc + 0.012))
            rii = len(lverts); lverts.append((xc + w, yy, zc + 0.012))
            left.append(li); right.append(rii)
        for k in range(nseg):
            lfaces.append((left[k], right[k], right[k + 1], left[k + 1]))
    return lverts, lfaces


for sign, tag in ((1, "N"), (-1, "S")):
    lv, lf = build_leaf(sign)
    leaf = new_mesh_object("Welwitschia_Leaf_%s" % tag, lv, lf, leaf_mat, smooth=True)
    try:
        sol = leaf.modifiers.new("Thickness", 'SOLIDIFY')
        sol.thickness = 0.02
        sol.offset = 0.0
    except Exception:
        pass
    created_objects.append(leaf)

# ---------------------------------------------------------------------------
# REPRODUCTIVE CONES on short branched peduncles rising from the crown margin
# ---------------------------------------------------------------------------
def add_cone_structure(loc, big):
    objs = []
    ped_h = 0.30 if big else 0.24
    ped_r = 0.028 if big else 0.020
    bpy.ops.mesh.primitive_cylinder_add(
        radius=ped_r, depth=ped_h, vertices=8,
        location=(loc.x, loc.y, loc.z + ped_h * 0.5))
    ped = bpy.context.object
    if ped is not None:
        if ped.data is not None:
            ped.data.materials.clear()
            ped.data.materials.append(stalk_mat)
        objs.append(ped)

    c_depth = 0.24 if big else 0.14
    c_r1 = 0.085 if big else 0.045
    c_r2 = 0.045 if big else 0.020
    ctop_z = loc.z + ped_h + c_depth * 0.5
    bpy.ops.mesh.primitive_cone_add(
        radius1=c_r1, radius2=c_r2, depth=c_depth, vertices=10,
        location=(loc.x, loc.y, ctop_z))
    cone = bpy.context.object
    if cone is not None:
        if cone.data is not None:
            cone.data.materials.clear()
            cone.data.materials.append(cone_mat)
        objs.append(cone)
    return objs


for i in range(12):
    th = random.uniform(0.0, TWO_PI)
    rr = random.uniform(0.18, 0.50)
    cx = rr * 1.12 * R * math.cos(th)
    cy = rr * 0.90 * R * math.sin(th)
    cz = 0.30
    created_objects.extend(add_cone_structure(Vector((cx, cy, cz)), big=(i % 3 == 0)))

# ---------------------------------------------------------------------------
# Group everything under one Empty so the specimen reads as a single organism
# ---------------------------------------------------------------------------
bpy.ops.object.select_all(action='DESELECT')
root_empty = bpy.data.objects.new("Welwitschia_mirabilis", None)
root_empty.empty_display_type = 'PLAIN_AXES'
root_empty.empty_display_size = 1.5
bpy.context.collection.objects.link(root_empty)

for ob in created_objects:
    if ob is None:
        continue
    try:
        ob.parent = root_empty
        ob.matrix_parent_inverse = root_empty.matrix_world.inverted()
    except Exception:
        pass

# leave the finished specimen in the scene (no saving, no clearing)