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

random.seed(20240925)

# ------------------------------------------------------------------
# Start from an empty scene so only the specimen remains
# ------------------------------------------------------------------
for _ob in list(bpy.data.objects):
    try:
        bpy.data.objects.remove(_ob, do_unlink=True)
    except Exception:
        pass

COLL = bpy.context.scene.collection


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


def nz(x, y=0.0, z=0.0):
    try:
        return noise.noise(Vector((x, y, z)))
    except Exception:
        return 0.5 * math.sin(x * 1.7 + y * 2.3 + z * 3.1)


def make_mat(name, color, rough=0.7):
    m = bpy.data.materials.new(name)
    try:
        m.diffuse_color = (color[0], color[1], color[2], 1.0)
    except Exception:
        pass
    try:
        m.use_nodes = True
        b = m.node_tree.nodes.get("Principled BSDF")
        if b is not None:
            try:
                b.inputs["Base Color"].default_value = (color[0], color[1], color[2], 1.0)
            except Exception:
                pass
            try:
                b.inputs["Roughness"].default_value = float(rough)
            except Exception:
                pass
    except Exception:
        pass
    return m


def loft(rings, closed=True, cap_start=None, cap_end=None):
    verts = []
    faces = []
    n = len(rings[0])
    for r in rings:
        verts.extend(r)
    rng = n if closed else n - 1
    for j in range(len(rings) - 1):
        for i in range(rng):
            a = j * n + i
            b = j * n + (i + 1) % n
            c = (j + 1) * n + (i + 1) % n
            d = (j + 1) * n + i
            faces.append((a, b, c, d))
    if cap_start is not None:
        idx = len(verts)
        verts.append(cap_start)
        for i in range(n):
            faces.append((idx, (i + 1) % n, i))
    if cap_end is not None:
        idx = len(verts)
        verts.append(cap_end)
        base = (len(rings) - 1) * n
        for i in range(n):
            faces.append((idx, base + i, base + (i + 1) % n))
    return verts, faces


def tube(points, radii, segs=8, caps=True):
    n = len(points)
    tans = []
    for i in range(n):
        if i == 0:
            t = points[1] - points[0]
        elif i == n - 1:
            t = points[-1] - points[-2]
        else:
            t = points[i + 1] - points[i - 1]
        if t.length < 1e-9:
            t = Vector((0, 0, 1))
        tans.append(t.normalized())
    u = tans[0].orthogonal().normalized()
    rings = []
    for i in range(n):
        t = tans[i]
        u = u - t * u.dot(t)
        if u.length < 1e-6:
            u = t.orthogonal()
        u.normalize()
        w = t.cross(u)
        ring = []
        for k in range(segs):
            a = 2.0 * math.pi * k / segs
            ring.append(points[i] + (u * math.cos(a) + w * math.sin(a)) * radii[i])
        rings.append(ring)
    if caps:
        return loft(rings, True, points[0] - tans[0] * radii[0] * 0.3,
                    points[-1] + tans[-1] * radii[-1] * 0.5)
    return loft(rings, True)


class Acc:
    def __init__(self):
        self.v = []
        self.f = []
        self.m = []

    def add(self, verts, faces, mi):
        o = len(self.v)
        self.v.extend(verts)
        self.f.extend([tuple(i + o for i in f) for f in faces])
        if isinstance(mi, list):
            self.m.extend(mi)
        else:
            self.m.extend([mi] * len(faces))


def build(name, verts, faces, mats, fmats=None, smooth=True):
    me = bpy.data.meshes.new(name)
    me.from_pydata([tuple(v) for v in verts], [], faces)
    for m in mats:
        me.materials.append(m)
    if fmats is not None and len(fmats) == len(me.polygons):
        try:
            me.polygons.foreach_set("material_index", fmats)
        except Exception:
            for p, i in zip(me.polygons, fmats):
                p.material_index = i
    try:
        me.validate()
    except Exception:
        pass
    try:
        bm = bmesh.new()
        bm.from_mesh(me)
        bmesh.ops.recalc_face_normals(bm, faces=bm.faces[:])
        bm.to_mesh(me)
        bm.free()
    except Exception:
        pass
    if smooth:
        try:
            me.polygons.foreach_set("use_smooth", [True] * len(me.polygons))
        except Exception:
            pass
    me.update()
    ob = bpy.data.objects.new(name, me)
    COLL.objects.link(ob)
    return ob


# ------------------------------------------------------------------
# Materials
# ------------------------------------------------------------------
M_BARK = make_mat("W_Bark_Cork", (0.14, 0.10, 0.07), 0.95)
M_CROWN = make_mat("W_Crown_Top", (0.23, 0.17, 0.11), 0.9)
M_ROOT = make_mat("W_Root", (0.36, 0.26, 0.16), 0.9)
M_LEAF_BASE = make_mat("W_Leaf_Meristem", (0.52, 0.60, 0.28), 0.55)
M_LEAF_GREEN = make_mat("W_Leaf_Green", (0.09, 0.24, 0.06), 0.5)
M_LEAF_OLIVE = make_mat("W_Leaf_Olive", (0.27, 0.24, 0.10), 0.7)
M_LEAF_DEAD = make_mat("W_Leaf_Dead", (0.34, 0.28, 0.20), 0.95)
M_PED = make_mat("W_Peduncle", (0.42, 0.40, 0.20), 0.7)
M_CONE_Y = make_mat("W_Cone_Young", (0.33, 0.48, 0.38), 0.5)
M_CONE_M = make_mat("W_Cone_Mature", (0.55, 0.13, 0.08), 0.55)

# ------------------------------------------------------------------
# Caudex (woody, obconical, two-lobed stem with concave crown)
# ------------------------------------------------------------------
A_R, B_R = 0.62, 0.34      # rim half-axes (long axis X)
Z_RIM = 0.26
Z_BOT = -0.40
NA = 128


def plan(theta):
    c, s = math.cos(theta), math.sin(theta)
    r = 1.0 / math.sqrt((c / A_R) ** 2 + (s / B_R) ** 2)
    d = math.atan2(abs(s), abs(c))           # angular distance from X axis
    r *= 1.0 - 0.10 * math.exp(-(d / 0.18) ** 2)   # notches at fissure ends
    return r


def fissure(y):
    return 0.085 * math.exp(-(y / 0.065) ** 2)


def side_f(z):
    t = max(0.0, (z - Z_BOT) / (Z_RIM - Z_BOT))
    return 0.35 + 0.65 * t ** 0.7


F_TOP0 = 0.93
rows = []
NS = 36
for k in range(NS + 1):
    z = Z_BOT + (Z_RIM - Z_BOT) * k / NS
    rows.append((side_f(z), z, 'side', 0.0))
for k in range(1, 6):
    a = k / 5.0 * math.pi / 2.0
    rows.append((1.0 - 0.07 * (1.0 - math.cos(a)), Z_RIM + 0.045 * math.sin(a), 'rim', k / 5.0))
NT = 16
for k in range(1, NT + 1):
    f = F_TOP0 * (1.0 - k / (NT + 0.6))
    rows.append((f, 0.0, 'top', 1.0))


def top_z(f, y):
    return (Z_RIM + 0.045 + 0.025 * (1.0 - f / F_TOP0)
            + 0.005 * math.sin(f * 60.0) - fissure(y))


caud_rings = []
for (f, zb, kind, fw) in rows:
    ring = []
    for i in range(NA):
        th = 2.0 * math.pi * i / NA
        r = plan(th) * f
        x, y = r * math.cos(th), r * math.sin(th)
        if kind == 'side':
            z = zb
            w_bark = 1.0 - sst(Z_RIM - 0.04, Z_RIM, z)
            c1 = abs(math.sin(th * 16.0 + 2.5 * nz(x * 1.5, y * 1.5, z * 3.0) + z * 4.0))
            crack_v = (1.0 - c1) ** 4
            c2 = abs(math.sin(z * 35.0 + 1.5 * nz(x * 2.0, y * 2.0, 3.0)))
            crack_h = (1.0 - c2) ** 6
            dr = (-0.022 * crack_v - 0.008 * crack_h
                  + 0.010 * nz(x * 6.0, y * 6.0, z * 6.0)) * w_bark
            # leaf-insertion grooves on both long sides
            g = math.exp(-((z - 0.19) / 0.03) ** 2)
            w_lat = 1.0 - sst(0.62 * A_R, 0.82 * A_R, abs(x))
            dr -= 0.035 * g * w_lat
            rr = r + dr
            x, y = rr * math.cos(th), rr * math.sin(th)
        elif kind == 'rim':
            z = zb - fissure(y) * fw + 0.004 * nz(x * 8, y * 8, 1.0)
        else:
            z = top_z(f, y) + 0.003 * nz(x * 10, y * 10, 2.0)
        ring.append(Vector((x, y, z)))
    caud_rings.append(ring)

cap_bot = Vector((0.0, 0.0, Z_BOT - 0.02))
cap_top = Vector((0.0, 0.0, top_z(0.0, 0.0)))
cv, cf = loft(caud_rings, True, cap_bot, cap_top)
cmat = []
for j in range(len(caud_rings) - 1):
    mi = 0 if j < NS + 5 else 1
    cmat.extend([mi] * NA)
cmat.extend([0] * NA)
cmat.extend([1] * NA)
caudex = build("Welwitschia_Caudex", cv, cf, [M_BARK, M_CROWN], cmat)

# ------------------------------------------------------------------
# Roots: deep tapering taproot with lateral roots and rootlets
# ------------------------------------------------------------------
root_acc = Acc()
tap_pts = []
tap_r = []
NTAP = 32
for k in range(NTAP):
    t = k / (NTAP - 1)
    z = -0.30 - 2.0 * t
    amp = 0.10 * t
    tap_pts.append(Vector((amp * nz(t * 2.5, 1.3, 0.0) * 2.0,
                           amp * nz(t * 2.5, 4.7, 0.0) * 2.0, z)))
    tap_r.append(0.125 * (1.0 - t) ** 1.2 + 0.004)
v, f = tube(tap_pts, tap_r, segs=16)
root_acc.add(v, f, 0)


def tap_at(t):
    fi = t * (NTAP - 1)
    i0 = int(math.floor(fi))
    i1 = min(i0 + 1, NTAP - 1)
    a = fi - i0
    return tap_pts[i0].lerp(tap_pts[i1], a), tap_r[i0] * (1 - a) + tap_r[i1] * a


def grow_root(start, direction, length, r0, npts, seed, down=0.35):
    pts = [start.copy()]
    rad = [r0]
    d = direction.normalized()
    step = length / (npts - 1)
    p = start.copy()
    for k in range(1, npts):
        t = k / (npts - 1)
        wob = Vector((nz(seed, t * 3.0, 0.0), nz(seed + 7.0, t * 3.0, 0.0), 0.0)) * 0.5
        d = (d + wob * 0.35 + Vector((0, 0, -down * t))).normalized()
        p = p + d * step
        pts.append(p.copy())
        rad.append(r0 * (1.0 - t) ** 1.3 + 0.0015)
    return pts, rad


for k in range(8):
    t = random.uniform(0.06, 0.55)
    p0, rt = tap_at(t)
    ang = k * 2.0 * math.pi / 8.0 + random.uniform(-0.3, 0.3)
    dirv = Vector((math.cos(ang), math.sin(ang), -random.uniform(0.1, 0.5)))
    start = p0 + Vector((math.cos(ang), math.sin(ang), 0.0)) * rt * 0.7
    L = random.uniform(0.6, 1.4)
    r0 = rt * random.uniform(0.25, 0.4)
    pts, rad = grow_root(start, dirv, L, r0, 16, random.uniform(0, 50))
    v, f = tube(pts, rad, segs=8)
    root_acc.add(v, f, 0)
    for q in range(3):
        idx = random.randint(3, 11)
        sp = pts[idx]
        a2 = ang + random.choice((-1, 1)) * random.uniform(0.6, 1.3)
        d2 = Vector((math.cos(a2), math.sin(a2), -random.uniform(0.2, 0.8)))
        pts2, rad2 = grow_root(sp, d2, random.uniform(0.15, 0.4),
                               rad[idx] * 0.5, 8, random.uniform(0, 50), 0.6)
        v, f = tube(pts2, rad2, segs=6)
        root_acc.add(v, f, 0)

roots = build("Welwitschia_Roots", root_acc.v, root_acc.f, [M_ROOT], root_acc.m)

# ------------------------------------------------------------------
# The two permanent strap leaves, torn into ribbons, lying on the ground
# ------------------------------------------------------------------
leaf_acc = Acc()
Z_LEAF0 = 0.19
f_ins = side_f(Z_LEAF0)
AE, BE = A_R * f_ins, B_R * f_ins
XS = 0.72 * AE


def leaf_z(s):
    return 0.012 + (Z_LEAF0 - 0.012) * math.exp(-s / 0.33) * (1.0 + 3.4 * s)


K = 5
for side in (1, -1):
    NSTRIP = 9
    cuts = [-1.0]
    for i in range(1, NSTRIP):
        cuts.append(-1.0 + 2.0 * i / NSTRIP + random.uniform(-0.06, 0.06))
    cuts.append(1.0)
    phi_side = side * math.pi / 2.0
    for si in range(NSTRIP):
        t0, t1 = cuts[si], cuts[si + 1]
        tc = 0.5 * (t0 + t1)
        bx = tc * XS
        q = max(0.05, 1.0 - (bx / AE) ** 2)
        by = side * (BE * math.sqrt(q) - 0.028)
        slope = BE * (abs(bx) / AE ** 2) / math.sqrt(q)
        w0 = (t1 - t0) * XS * math.sqrt(1.0 + slope * slope) * 1.03
        phin = math.atan2(by / BE ** 2, bx / AE ** 2)
        phi0 = 0.55 * phin + 0.45 * phi_side
        fan = (phi0 - phi_side) * 0.9 + random.uniform(-0.2, 0.2)
        L = random.uniform(1.6, 2.8)
        s_split = random.uniform(0.25, 0.6)
        seed = random.uniform(0.0, 100.0)
        roll_amp = random.uniform(0.2, 0.9)
        d1 = random.uniform(0.45, 0.62)
        d2 = d1 + random.uniform(0.1, 0.18)
        ds = 0.025
        steps = int(L / ds)
        pos = Vector((bx, by))
        sections = []
        for j in range(steps + 1):
            s = j * ds
            phi = (phi0 + fan * sst(s_split, s_split + 1.3, s)
                   + 0.45 * nz(s * 0.7, seed, 0.0) * sst(s_split, s_split + 1.0, s))
            taper = 1.0 - 0.85 * sst(L - 0.35, L, s)
            fray = 1.0 + 0.12 * nz(s * 6.0, seed, 3.0) * sst(0.6 * L, L, s)
            w = w0 * (1.0 + 0.18 * min(s, 1.5)) * taper * fray
            curl = 0.12 + 0.45 * sst(0.6 * L, L, s)
            roll = roll_amp * nz(s * 0.6 + 3.0, seed + 5.0, 1.0) * sst(s_split, s_split + 0.8, s)
            z = (leaf_z(s) + 0.01 * nz(s * 1.3, seed, 7.0) * sst(0.8, 1.4, s)
                 + 0.0015 * (si % 3))
            z += abs(math.sin(roll)) * w * 0.5
            lat = Vector((-math.sin(phi), math.cos(phi), 0.0))
            up = Vector((0.0, 0.0, 1.0))
            lat_r = lat * math.cos(roll) + up * math.sin(roll)
            up_r = up * math.cos(roll) - lat * math.sin(roll)
            C = Vector((pos.x, pos.y, z))
            sec = []
            for k in range(K):
                u = -1.0 + 2.0 * k / (K - 1)
                sec.append(C + lat_r * (u * w * 0.5) + up_r * (curl * u * u * w * 0.5))
            sections.append((sec, s))
            pos = pos + Vector((math.cos(phi), math.sin(phi))) * ds
        verts = []
        faces = []
        fm = []
        for (sec, s) in sections:
            verts.extend(sec)
        for j in range(len(sections) - 1):
            fr = sections[j][1] / L
            if fr < 0.07:
                mi = 0
            elif fr < d1:
                mi = 1
            elif fr < d2:
                mi = 2
            else:
                mi = 3
            for k in range(K - 1):
                a = j * K + k
                faces.append((a, a + 1, a + K + 1, a + K))
                fm.append(mi)
        leaf_acc.add(verts, faces, fm)

leaves = build("Welwitschia_Leaves", leaf_acc.v, leaf_acc.f,
               [M_LEAF_BASE, M_LEAF_GREEN, M_LEAF_OLIVE, M_LEAF_DEAD], leaf_acc.m)
try:
    sol = leaves.modifiers.new("LeafThickness", 'SOLIDIFY')
    sol.thickness = 0.004
    sol.offset = 0.0
except Exception:
    pass

# ------------------------------------------------------------------
# Female strobili on dichasially branched peduncles from the crown lobes
# ------------------------------------------------------------------
cone_acc = Acc()


def make_cone(base, d, length, radius, mi):
    d = d.normalized()
    u = d.orthogonal().normalized()
    w = d.cross(u)
    NV, NU = 22, 24
    rings = []
    for j in range(1, NV):
        v = j / NV
        prof = radius * (math.sin(math.pi * v ** 0.9)) ** 0.75
        whorl = v * 9.0
        parity = int(math.floor(whorl)) % 2
        saw = whorl - math.floor(whorl)
        ring = []
        for i in range(NU):
            a = 2.0 * math.pi * i / NU
            q = (a / (math.pi / 2.0) + 0.5 * parity) % 1.0   # decussate 4-row bracts
            bump = math.sin(math.pi * q) * saw ** 1.5
            r = prof * (1.0 + 0.16 * bump)
            ring.append(base + d * (v * length) + (u * math.cos(a) + w * math.sin(a)) * r)
        rings.append(ring)
    v, f = loft(rings, True, base.copy(), base + d * length)
    cone_acc.add(v, f, mi)


def peduncle(p, d, length, depth, rot):
    d = d.normalized()
    q = p + d * length
    bend = d.orthogonal().normalized() * (length * 0.06 * random.uniform(-1, 1))
    mid = p + d * (length * 0.5) + bend
    r0 = 0.0028 + 0.0012 * depth
    v, f = tube([p, mid, q], [r0, r0 * 0.95, r0 * 0.88], segs=8)
    cone_acc.add(v, f, 0)
    if depth == 0:
        clen = random.uniform(0.04, 0.065)
        mi = 1 if random.random() < 0.55 else 2
        make_cone(q - d * 0.003, d, clen, clen * 0.38, mi)
        return
    ax0 = d.orthogonal().normalized()
    axis = Matrix.Rotation(rot, 3, d) @ ax0
    ang = random.uniform(0.45, 0.7)
    for sgn in (-1.0, 1.0):
        nd = Matrix.Rotation(sgn * ang, 3, axis) @ d
        nd = (nd + Vector((0.0, 0.0, 0.25))).normalized()
        peduncle(q, nd, length * random.uniform(0.65, 0.8), depth - 1, rot + math.pi / 2.0)


for side in (1, -1):
    for k in range(5):
        x = -0.36 + k * 0.18 + random.uniform(-0.03, 0.03)
        yy = side * B_R * math.sqrt(max(0.05, 1.0 - (x / A_R) ** 2)) * random.uniform(0.66, 0.76)
        fr = math.sqrt((x / A_R) ** 2 + (yy / B_R) ** 2)
        z = top_z(min(fr, F_TOP0), yy) - 0.008
        d = Vector((random.uniform(-0.15, 0.15), side * 0.35, 1.0)).normalized()
        peduncle(Vector((x, yy, z)), d, random.uniform(0.05, 0.08),
                 random.choice((2, 2, 3)), random.uniform(0, math.pi))

cones = build("Welwitschia_Strobili", cone_acc.v, cone_acc.f,
              [M_PED, M_CONE_Y, M_CONE_M], cone_acc.m)

# ------------------------------------------------------------------
# Group everything under one specimen root
# ------------------------------------------------------------------
root_empty = bpy.data.objects.new("Welwitschia_mirabilis", None)
COLL.objects.link(root_empty)
for name in ("Welwitschia_Caudex", "Welwitschia_Roots",
             "Welwitschia_Leaves", "Welwitschia_Strobili"):
    ob = bpy.data.objects.get(name)
    if ob is not None:
        ob.parent = root_empty

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