import bpy
import bmesh
import math
import random
from mathutils import Vector, Matrix, Quaternion, noise

random.seed(20240925)

# ---------------------------------------------------------------------------
# Remove default mesh objects so only the specimen remains
# ---------------------------------------------------------------------------
for _ob in list(bpy.data.objects):
    if _ob.type == 'MESH':
        try:
            bpy.data.objects.remove(_ob, do_unlink=True)
        except Exception:
            pass

coll = bpy.data.collections.new("Welwitschia_mirabilis")
bpy.context.scene.collection.children.link(coll)

root = bpy.data.objects.new("Welwitschia_mirabilis", None)
coll.objects.link(root)


# ---------------------------------------------------------------------------
# Utilities
# ---------------------------------------------------------------------------
def smoothstep(a, b, x):
    if b == a:
        return 1.0 if x >= b else 0.0
    t = max(0.0, min(1.0, (x - a) / (b - a)))
    return t * t * (3.0 - 2.0 * t)


def lerp(a, b, t):
    return a + (b - a) * t


def lerp3(c1, c2, t):
    return (lerp(c1[0], c2[0], t), lerp(c1[1], c2[1], t), lerp(c1[2], c2[2], t))


def make_mat(name, color, rough=0.7, vcol=None, bump=0.0, bump_scale=40.0, sss=0.0):
    mat = bpy.data.materials.new(name)
    try:
        mat.diffuse_color = (color[0], color[1], color[2], 1.0)
    except Exception:
        pass
    try:
        mat.use_nodes = True
        nt = mat.node_tree
        bsdf = nt.nodes.get("Principled BSDF")
        if bsdf is None:
            for n in nt.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
            if vcol:
                try:
                    an = nt.nodes.new('ShaderNodeAttribute')
                    an.attribute_name = vcol
                    nt.links.new(an.outputs['Color'], bsdf.inputs['Base Color'])
                except Exception:
                    pass
            if sss > 0.0:
                try:
                    bsdf.inputs['Subsurface Weight'].default_value = float(sss)
                except Exception:
                    pass
            if bump > 0.0:
                try:
                    tex = nt.nodes.new('ShaderNodeTexNoise')
                    tex.inputs['Scale'].default_value = float(bump_scale)
                    bp = nt.nodes.new('ShaderNodeBump')
                    bp.inputs['Strength'].default_value = float(bump)
                    nt.links.new(tex.outputs['Fac'], bp.inputs['Height'])
                    nt.links.new(bp.outputs['Normal'], bsdf.inputs['Normal'])
                except Exception:
                    pass
    except Exception:
        pass
    return mat


def finish(name, bm, mat, colors=None, smooth=True):
    if smooth:
        for f in bm.faces:
            f.smooth = True
    me = bpy.data.meshes.new(name + "_mesh")
    bm.to_mesh(me)
    bm.free()
    if colors:
        try:
            attr = me.color_attributes.new(name="Col", type='FLOAT_COLOR', domain='POINT')
            flat = []
            for c in colors:
                flat.extend((c[0], c[1], c[2], 1.0))
            if len(flat) == len(attr.data) * 4:
                attr.data.foreach_set("color", flat)
        except Exception as e:
            print("color attribute warning:", e)
    ob = bpy.data.objects.new(name, me)
    coll.objects.link(ob)
    if mat is not None:
        me.materials.append(mat)
    ob.parent = root
    return ob


def tube(bm, pts, radii, segs=10, cap_start=True, tip=True):
    """Sweep a circular cross-section along a polyline (parallel transport frames)."""
    n = len(pts)
    tangents = []
    for i in range(n):
        a = pts[max(0, i - 1)]
        b = pts[min(n - 1, i + 1)]
        t = b - a
        if t.length < 1e-9:
            t = Vector((0.0, 0.0, 1.0))
        tangents.append(t.normalized())
    ref = Vector((0.0, 0.0, 1.0)) if abs(tangents[0].z) < 0.9 else Vector((1.0, 0.0, 0.0))
    nrm = tangents[0].cross(ref)
    if nrm.length < 1e-6:
        nrm = tangents[0].orthogonal()
    nrm.normalize()
    rings = []
    for i in range(n):
        t = tangents[i]
        nrm = nrm - t * nrm.dot(t)
        if nrm.length < 1e-6:
            nrm = t.orthogonal()
        nrm.normalize()
        bn = t.cross(nrm).normalized()
        ring = []
        for k in range(segs):
            ang = 2.0 * math.pi * k / segs
            ring.append(bm.verts.new(pts[i] + (nrm * math.cos(ang) + bn * math.sin(ang)) * radii[i]))
        rings.append(ring)
    for i in range(n - 1):
        for k in range(segs):
            k2 = (k + 1) % segs
            bm.faces.new((rings[i][k], rings[i][k2], rings[i + 1][k2], rings[i + 1][k]))
    if cap_start:
        c = bm.verts.new(pts[0] - tangents[0] * radii[0] * 0.3)
        for k in range(segs):
            bm.faces.new((rings[0][(k + 1) % segs], rings[0][k], c))
    if tip:
        c = bm.verts.new(pts[-1] + tangents[-1] * max(radii[-1], 0.002))
        for k in range(segs):
            bm.faces.new((rings[-1][k], rings[-1][(k + 1) % segs], c))


# ---------------------------------------------------------------------------
# Materials
# ---------------------------------------------------------------------------
mat_bark = make_mat("Welw_Bark", (0.16, 0.12, 0.09), rough=0.95, bump=0.6, bump_scale=35.0)
mat_root = make_mat("Welw_Root", (0.30, 0.22, 0.15), rough=0.9, bump=0.3, bump_scale=25.0)
mat_leaf = make_mat("Welw_Leaf", (0.20, 0.28, 0.10), rough=0.55, vcol="Col", bump=0.08,
                    bump_scale=120.0, sss=0.03)
mat_stalk = make_mat("Welw_Peduncle", (0.38, 0.36, 0.20), rough=0.6)
mat_cone_mature = make_mat("Welw_FemaleCone_Mature", (0.55, 0.22, 0.15), rough=0.5, bump=0.1)
mat_cone_young = make_mat("Welw_FemaleCone_Young", (0.30, 0.42, 0.34), rough=0.45, bump=0.1)

# ---------------------------------------------------------------------------
# Woody crown (short, broad, two-lobed, concave obconical stem)
# ---------------------------------------------------------------------------
RX, RY = 0.34, 0.56   # leaves emerge from the long (Y) margins, extend along +/-X

prof = [(0.80, -0.12), (0.83, -0.05), (0.87, 0.03), (0.91, 0.11), (0.95, 0.19),
        (0.985, 0.26), (1.0, 0.305), (0.975, 0.33), (0.93, 0.328), (0.86, 0.312),
        (0.76, 0.292), (0.63, 0.272), (0.49, 0.258), (0.35, 0.248), (0.21, 0.242),
        (0.09, 0.238)]
NT = 96
CLEFT_DEPTH = 0.075


def cleft(x):
    return CLEFT_DEPTH * math.exp(-(x / 0.07) ** 2)


bm = bmesh.new()
rings = []
for (s, z) in prof:
    wtop = smoothstep(0.18, 0.30, z)
    ring = []
    for i in range(NT):
        th = 2.0 * math.pi * i / NT
        c, sn = math.cos(th), math.sin(th)
        rn = 1.0
        rn += 0.035 * noise.noise(Vector((c * 2.5, sn * 2.5, z * 4.0)))
        rn += 0.018 * noise.noise(Vector((c * 14.0, sn * 14.0, z * 0.8))) * (1.0 - wtop)
        x = s * RX * c * rn
        y = s * RY * sn * rn
        zz = z - cleft(x) * wtop
        zz += 0.014 * noise.noise(Vector((x * 11.0, y * 11.0, 5.0))) * wtop
        zz += 0.006 * noise.noise(Vector((x * 40.0, y * 40.0, 2.0))) * wtop
        ring.append(bm.verts.new((x, y, zz)))
    rings.append(ring)

for k in range(len(rings) - 1):
    for i in range(NT):
        i2 = (i + 1) % NT
        bm.faces.new((rings[k][i], rings[k][i2], rings[k + 1][i2], rings[k + 1][i]))
top_c = bm.verts.new((0.0, 0.0, 0.232 - CLEFT_DEPTH))
for i in range(NT):
    bm.faces.new((rings[-1][i], rings[-1][(i + 1) % NT], top_c))
bot_c = bm.verts.new((0.0, 0.0, -0.13))
for i in range(NT):
    bm.faces.new((rings[0][(i + 1) % NT], rings[0][i], bot_c))
bmesh.ops.recalc_face_normals(bm, faces=bm.faces)
crown = finish("Crown_WoodyStem", bm, mat_bark)
try:
    m = crown.modifiers.new("Subsurf", 'SUBSURF')
    m.levels = 1
    m.render_levels = 2
except Exception:
    pass

# ---------------------------------------------------------------------------
# Taproot with lateral roots
# ---------------------------------------------------------------------------
bm = bmesh.new()
pts, radii = [], []
NTR = 26
for i in range(NTR):
    t = i / (NTR - 1)
    z = -0.05 - t * 1.45
    x = 0.05 * math.sin(t * 5.0) + 0.02 * noise.noise(Vector((t * 6.0, 0.3, 0.1)))
    y = 0.04 * math.sin(t * 3.7 + 1.0)
    pts.append(Vector((x, y, z)))
    radii.append(0.25 * (1.0 - t) ** 1.7 + 0.006)
tube(bm, pts, radii, segs=16, cap_start=True, tip=True)

for li in range(6):
    t0 = random.uniform(0.18, 0.65)
    idx = int(t0 * (NTR - 1))
    start = pts[idx].copy()
    ang = li * (2.0 * math.pi / 6.0) + random.uniform(-0.4, 0.4)
    horiz = Vector((math.cos(ang), math.sin(ang), 0.0))
    L = random.uniform(0.45, 0.9)
    r0 = radii[idx] * 0.35
    lp, lr = [], []
    n = 14
    for j in range(n):
        f = j / (n - 1)
        p = start + horiz * (f * L + radii[idx] * 0.5) + Vector((0.0, 0.0, -f * L * 0.55 - 0.3 * f * f * L))
        p += Vector((0.03 * math.sin(f * 9 + li), 0.03 * math.cos(f * 7 + li), 0.0)) * f
        lp.append(p)
        lr.append(r0 * (1.0 - f) ** 1.3 + 0.002)
    tube(bm, lp, lr, segs=8, cap_start=True, tip=True)
bmesh.ops.recalc_face_normals(bm, faces=bm.faces)
finish("Taproot", bm, mat_root)

# ---------------------------------------------------------------------------
# The two strap-shaped leaves (split, frayed, curled, dead tips)
# ---------------------------------------------------------------------------
ZG = 0.012            # leaf resting height above ground (ground = z 0)
Z0 = 0.30             # leaf insertion height in crown groove
W0 = 0.90             # leaf width at base


def build_path():
    ds = 0.01
    d, z = 0.0, Z0
    phi = math.radians(22.0)
    D, Zs = [d], [z]
    for _ in range(500):
        phi_lim = -math.radians(38.0) * min(1.0, max(0.0, (z - ZG) / 0.10))
        phi = max(phi - math.radians(120.0) * ds, phi_lim)
        d += math.cos(phi) * ds
        z = max(ZG, z + math.sin(phi) * ds)
        D.append(d)
        Zs.append(z)
    return D, Zs, ds


PD, PZ, PDS = build_path()


def path_at(u):
    idx = max(0.0, u) / PDS
    i = int(min(idx, len(PD) - 2))
    f = min(1.0, idx - i)
    return lerp(PD[i], PD[i + 1], f), lerp(PZ[i], PZ[i + 1], f)


def xb(y):
    return 0.95 * RX * math.sqrt(max(0.04, 1.0 - (y / RY) ** 2))


XB0 = xb(0.0)

C_BASE = (0.36, 0.42, 0.20)
C_GREEN = (0.10, 0.24, 0.06)
C_OLIVE = (0.24, 0.27, 0.12)
C_TAN = (0.42, 0.33, 0.20)
C_DEAD = (0.22, 0.17, 0.13)


def build_leaf(sgn, name):
    bm = bmesh.new()
    colors = []
    N = random.randint(7, 9)
    ws = [random.uniform(0.6, 1.4) for _ in range(N)]
    tot = sum(ws)
    edges = [-W0 / 2.0]
    for w in ws:
        edges.append(edges[-1] + w / tot * W0)
    NC = 6
    for i in range(N):
        yc0 = 0.5 * (edges[i] + edges[i + 1])
        w0 = edges[i + 1] - edges[i]
        us = random.uniform(0.3, 1.1)                  # where this strip splits free
        Li = random.uniform(1.7, 2.8)                  # strip length (tips abraded)
        tau = random.choice((-1.0, 1.0)) * random.uniform(0.3, 1.4)
        drift_k = random.uniform(-0.08, 0.08) + 0.07 * (yc0 / (W0 / 2.0))
        ph = random.uniform(0.0, 6.28)
        wave_amp = random.uniform(0.01, 0.05)
        wave_fr = random.uniform(3.0, 6.0)
        wave_ph = random.uniform(0.0, 6.28)
        tip_curl = random.uniform(0.03, 0.14)
        gshift = random.uniform(-0.15, 0.15)
        layer = 0.003 * (i % 3)
        NR = max(24, int(Li / 0.025))
        cuts = [random.uniform(0.0, 0.18) for _ in range(NC + 1)]
        grid = []
        for r in range(NR + 1):
            fr = r / NR
            row = []
            for j in range(NC + 1):
                v = j / NC - 0.5
                u = fr * Li - cuts[j] * fr ** 8
                F = 1.0 + 0.5 * smoothstep(0.25, 1.6, u)
                sp = max(0.0, u - us)
                drift = drift_k * sp + 0.035 * math.sin(2.1 * u + ph) * min(1.0, sp * 2.0)
                yc = yc0 * F + drift
                w = w0 * F
                w *= (1.0 - 0.3 * smoothstep(us, us + 0.8, u))
                w *= (1.0 - 0.6 * smoothstep(Li - 0.35, Li, u))
                dist, zc = path_at(u)
                y0col = yc0 + v * w0
                xoff = lerp(xb(y0col), lerp(xb(y0col), XB0, 0.6), smoothstep(0.0, 1.0, u))
                x = sgn * (xoff - 0.03 + dist)
                tw = tau * smoothstep(us, Li, u)
                lv = v * w
                y = yc + lv * math.cos(tw)
                z = zc + lv * math.sin(tw) + 0.5 * w * abs(math.sin(tw))
                z += (0.05 + 0.25 * smoothstep(1.0, Li, u)) * w * (2.0 * v) ** 2
                z += wave_amp * max(0.0, math.sin(u * wave_fr + wave_ph)) ** 2 * smoothstep(0.9, 1.3, u)
                z += tip_curl * smoothstep(Li - 0.4, Li, u) ** 2
                z += layer * smoothstep(us, us + 0.3, u)
                z += 0.004 * noise.noise(Vector((x * 6.0, y * 6.0, 0.7)))
                z = max(ZG, z)
                row.append(bm.verts.new((x, y, z)))
                # colour: pale meristematic base -> green -> olive -> tan -> dead tip
                col = lerp3(C_BASE, C_GREEN, smoothstep(0.0, 0.15, u))
                col = lerp3(col, C_OLIVE, smoothstep(0.35 + gshift, 0.95 + gshift, u))
                col = lerp3(col, C_TAN, smoothstep(0.95 + gshift, 1.6 + gshift, u))
                col = lerp3(col, C_DEAD, smoothstep(Li - 0.6, Li, u))
                jit = random.uniform(-0.015, 0.015)
                colors.append((max(0.0, col[0] + jit), max(0.0, col[1] + jit), max(0.0, col[2] + jit)))
            grid.append(row)
        for r in range(NR):
            for j in range(NC):
                bm.faces.new((grid[r][j], grid[r][j + 1], grid[r + 1][j + 1], grid[r + 1][j]))
    ob = finish(name, bm, mat_leaf, colors=colors)
    try:
        m = ob.modifiers.new("Solidify", 'SOLIDIFY')
        m.thickness = 0.007
        m.offset = 0.0
    except Exception:
        pass
    return ob


build_leaf(1.0, "Leaf_A")
build_leaf(-1.0, "Leaf_B")

# ---------------------------------------------------------------------------
# Female strobili on dichasially branched peduncles arising from the crown rim
# ---------------------------------------------------------------------------
bm_stalk = bmesh.new()
bm_cone_m = bmesh.new()
bm_cone_y = bmesh.new()


def add_female_cone(bm, base, direction, h, r):
    d = direction.normalized()
    R = d.to_track_quat('Z', 'Y').to_matrix().to_4x4()
    Tc = Matrix.Translation(base + d * (h * 0.5))
    bmesh.ops.create_icosphere(bm, subdivisions=2, radius=1.0,
                               matrix=Tc @ R @ Matrix.Diagonal((r * 0.68, r * 0.68, h * 0.5, 1.0)))
    npairs = 9
    for j in range(npairs):
        t = (j + 0.5) / npairs
        zl = -h * 0.5 + t * h
        pr = math.sqrt(max(0.0, 1.0 - (zl / (h * 0.5)) ** 2))
        rr = r * max(0.35, pr)
        for k in range(2):  # decussate bract pairs -> four ranks
            a = j * math.pi * 0.5 + k * math.pi + random.uniform(-0.08, 0.08)
            M = (Tc @ R @ Matrix.Translation((0.0, 0.0, zl)) @ Matrix.Rotation(a, 4, 'Z')
                 @ Matrix.Translation((rr * 0.72, 0.0, 0.0)) @ Matrix.Rotation(0.3, 4, 'Y')
                 @ Matrix.Diagonal((r * 0.2, rr * 0.95, h * 0.17, 1.0)))
            bmesh.ops.create_icosphere(bm, subdivisions=1, radius=1.0, matrix=M)


def dichasium(p, d, length, rad, depth, plane_angle):
    bend = Vector((random.uniform(-1, 1), random.uniform(-1, 1), 0.0)) * length * 0.05
    end = p + d * length
    mid = (p + end) * 0.5 + bend
    tube(bm_stalk, [p, mid, end], [rad, rad * 0.93, rad * 0.86], segs=8, cap_start=True, tip=False)
    if depth == 0:
        target = bm_cone_m if random.random() < 0.6 else bm_cone_y
        add_female_cone(target, end - d * 0.004, d, random.uniform(0.05, 0.075), random.uniform(0.017, 0.023))
        return
    # small scale-like bracts at the branching node
    bmesh.ops.create_icosphere(bm_stalk, subdivisions=1, radius=1.0,
                               matrix=Matrix.Translation(end) @ Matrix.Diagonal((rad * 1.5, rad * 1.5, rad * 1.2, 1.0)))
    perp = d.orthogonal().normalized()
    perp.rotate(Quaternion(d, plane_angle))
    for s in (-1.0, 1.0):
        nd = d.copy()
        nd.rotate(Quaternion(perp, s * random.uniform(0.45, 0.7)))
        if nd.z < 0.25:
            nd.z = 0.25
        nd.normalize()
        dichasium(end, nd, length * 0.75, rad * 0.8, depth - 1, plane_angle + math.pi * 0.5)


for sgn in (1.0, -1.0):
    for yy in (-0.32, -0.11, 0.11, 0.32):
        y = yy + random.uniform(-0.04, 0.04)
        x = sgn * xb(y) * random.uniform(0.72, 0.82)
        p = Vector((x, y, 0.27))
        d = Vector((sgn * random.uniform(0.2, 0.5), random.uniform(-0.2, 0.2), 1.0)).normalized()
        dichasium(p, d, random.uniform(0.07, 0.10), 0.006, random.choice((1, 2)), random.uniform(0.0, math.pi))

for _bm in (bm_stalk, bm_cone_m, bm_cone_y):
    try:
        bmesh.ops.recalc_face_normals(_bm, faces=_bm.faces)
    except Exception:
        pass

finish("Strobilus_Peduncles", bm_stalk, mat_stalk)
finish("Female_Strobili_Mature", bm_cone_m, mat_cone_mature)
finish("Female_Strobili_Young", bm_cone_y, mat_cone_young)

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