import bpy
import bmesh
import math
import random
from math import pi, cos, sin

random.seed(7)

# ---------------------------------------------------------------------------
# Utility: start from an empty scene (delete existing objects safely)
# ---------------------------------------------------------------------------
bpy.ops.object.select_all(action='DESELECT')
for ob in list(bpy.data.objects):
    bpy.data.objects.remove(ob, do_unlink=True)


# ---------------------------------------------------------------------------
# Material helper (minimal + defensive)
# ---------------------------------------------------------------------------
def make_material(name, color, rough=0.6, subsurface=0.0):
    mat = bpy.data.materials.new(name)
    mat.use_nodes = True
    nodes = mat.node_tree.nodes
    bsdf = None
    for n in nodes:
        if n.type == 'BSDF_PRINCIPLED':
            bsdf = n
            break
    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
        # Blender 4.x renamed subsurface sockets -> guard everything
        try:
            bsdf.inputs["Subsurface Weight"].default_value = float(subsurface)
        except Exception:
            pass
        try:
            bsdf.inputs["Subsurface Radius"].default_value = (color[0], color[1], color[2])
        except Exception:
            pass
    return mat


mat_cap = make_material("Cap_Red", (0.75, 0.03, 0.02), rough=0.35, subsurface=0.25)
mat_white = make_material("Veil_White", (0.93, 0.90, 0.82), rough=0.55, subsurface=0.35)
mat_stipe = make_material("Stipe_White", (0.95, 0.93, 0.86), rough=0.6, subsurface=0.4)
mat_gills = make_material("Gills_Cream", (0.90, 0.87, 0.78), rough=0.7, subsurface=0.3)


# ---------------------------------------------------------------------------
# Mesh helpers
# ---------------------------------------------------------------------------
def obj_from_bmesh(name, bm, mat=None):
    mesh = bpy.data.meshes.new(name)
    for f in bm.faces:
        f.smooth = True
    bm.to_mesh(mesh)
    bm.free()
    ob = bpy.data.objects.new(name, mesh)
    bpy.context.collection.objects.link(ob)
    if mat is not None:
        ob.data.materials.append(mat)
    return ob


def revolve(name, profile, mat=None, segments=64, noise=0.0):
    """Revolve a 2D profile [(r, z), ...] around the Z axis."""
    bm = bmesh.new()
    verts = [bm.verts.new((r, 0.0, z)) for (r, z) in profile]
    edges = [bm.edges.new((verts[i], verts[i + 1])) for i in range(len(verts) - 1)]
    geom = verts + edges
    bmesh.ops.spin(bm, geom=geom, cent=(0, 0, 0), axis=(0, 0, 1),
                   angle=2 * pi, steps=segments, use_duplicate=False)
    bmesh.ops.remove_doubles(bm, verts=bm.verts, dist=1e-4)
    if noise > 0.0:
        for v in bm.verts:
            v.co.z += random.uniform(-noise, noise)
            v.co.x += random.uniform(-noise, noise)
            v.co.y += random.uniform(-noise, noise)
    bmesh.ops.recalc_face_normals(bm, faces=bm.faces)
    return obj_from_bmesh(name, bm, mat)


def add_blob(bm, center, radius, squash_z=0.5, subdivisions=1):
    """Add a small flattened icosphere (a veil wart) into an existing bmesh."""
    res = bmesh.ops.create_icosphere(bm, subdivisions=subdivisions, radius=radius)
    vlist = res['verts']
    bmesh.ops.scale(bm, vec=(1.0, 1.0, squash_z), verts=vlist)
    bmesh.ops.translate(bm, vec=center, verts=vlist)


# ---------------------------------------------------------------------------
# Cap surface profile + height lookup
# ---------------------------------------------------------------------------
CAP_PROFILE = [
    (0.0, 3.40),
    (0.40, 3.36),
    (0.80, 3.22),
    (1.15, 2.98),
    (1.40, 2.72),
    (1.50, 2.58),
    (1.52, 2.50),
    (1.40, 2.48),
]


def cap_z(r):
    pts = CAP_PROFILE[:6]  # only the upper (outer) surface
    if r <= pts[0][0]:
        return pts[0][1]
    if r >= pts[-1][0]:
        return pts[-1][1]
    for i in range(len(pts) - 1):
        r0, z0 = pts[i]
        r1, z1 = pts[i + 1]
        if r0 <= r <= r1:
            t = (r - r0) / (r1 - r0)
            return z0 + t * (z1 - z0)
    return pts[-1][1]


# ---------------------------------------------------------------------------
# 1. Cap (pileus) — red convex cap with slightly down-turned margin
# ---------------------------------------------------------------------------
cap = revolve("Cap", CAP_PROFILE, mat_cap, segments=80, noise=0.006)

# ---------------------------------------------------------------------------
# 2. White warts (remnants of the universal veil) on top of the cap
# ---------------------------------------------------------------------------
bm_w = bmesh.new()
n_warts = 55
for i in range(n_warts):
    r = math.sqrt(random.uniform(0.0, 1.0)) * 1.42
    ang = random.uniform(0.0, 2 * pi)
    x = r * cos(ang)
    y = r * sin(ang)
    z = cap_z(r) + 0.01
    size = max(0.03, random.uniform(0.05, 0.11) * (1.0 - 0.35 * (r / 1.45)))
    add_blob(bm_w, (x, y, z), size, squash_z=0.45)
bmesh.ops.recalc_face_normals(bm_w, faces=bm_w.faces)
warts = obj_from_bmesh("Warts", bm_w, mat_white)

# ---------------------------------------------------------------------------
# 3. Gills (lamellae) — thin radial blades beneath the cap
# ---------------------------------------------------------------------------
bm_g = bmesh.new()
n_gills = 72
inner_r = 0.33
outer_r = 1.44
for i in range(n_gills):
    ang = 2 * pi * i / n_gills
    ca, sa = cos(ang), sin(ang)
    ti = (inner_r * ca, inner_r * sa, 2.66)   # top inner
    to = (outer_r * ca, outer_r * sa, 2.50)   # top outer
    bo = (outer_r * ca, outer_r * sa, 2.36)   # bottom (free) outer
    bi = (inner_r * ca, inner_r * sa, 2.52)   # bottom inner
    v1 = bm_g.verts.new(ti)
    v2 = bm_g.verts.new(to)
    v3 = bm_g.verts.new(bo)
    v4 = bm_g.verts.new(bi)
    bm_g.faces.new((v1, v2, v3, v4))
bmesh.ops.recalc_face_normals(bm_g, faces=bm_g.faces)
gills = obj_from_bmesh("Gills", bm_g, mat_gills)

# ---------------------------------------------------------------------------
# 4. Stipe (stem) — tapered, with a swollen bulbous base
# ---------------------------------------------------------------------------
STIPE_PROFILE = [
    (0.00, 0.00),
    (0.50, 0.02),
    (0.62, 0.16),
    (0.60, 0.42),
    (0.45, 0.72),
    (0.36, 1.20),
    (0.32, 1.80),
    (0.31, 2.30),
    (0.34, 2.70),
    (0.30, 2.85),
    (0.00, 2.88),
]
stipe = revolve("Stipe", STIPE_PROFILE, mat_stipe, segments=48, noise=0.004)

# ---------------------------------------------------------------------------
# 5. Ring / annulus — skirt-like membrane around the upper stipe
# ---------------------------------------------------------------------------
RING_PROFILE = [
    (0.34, 2.42),
    (0.46, 2.32),
    (0.53, 2.16),
    (0.51, 2.05),
    (0.42, 2.04),
]
ring = revolve("Ring", RING_PROFILE, mat_white, segments=48)

# ---------------------------------------------------------------------------
# 6. Volva — concentric rings of scaly remnants on the bulbous base
# ---------------------------------------------------------------------------
bm_v = bmesh.new()
for band_z, band_r, count in [(0.12, 0.60, 22), (0.30, 0.61, 20), (0.48, 0.53, 16)]:
    for j in range(count):
        ang = 2 * pi * j / count + random.uniform(-0.1, 0.1)
        x = band_r * cos(ang)
        y = band_r * sin(ang)
        z = band_z + random.uniform(-0.02, 0.02)
        size = random.uniform(0.05, 0.09)
        add_blob(bm_v, (x, y, z), size, squash_z=0.6)
bmesh.ops.recalc_face_normals(bm_v, faces=bm_v.faces)
volva = obj_from_bmesh("Volva", bm_v, mat_white)

# ---------------------------------------------------------------------------
# Finalise: name the whole specimen root, leave everything in the scene
# ---------------------------------------------------------------------------
bpy.ops.object.select_all(action='DESELECT')