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

random.seed(7)
Z = Vector((0.0, 0.0, 1.0))

# ---------------------------------------------------------------------------
# Clean start (scene is intentionally cleared only at the beginning)
# ---------------------------------------------------------------------------
for _ob in list(bpy.data.objects):
    try:
        bpy.data.objects.remove(_ob, do_unlink=True)
    except Exception:
        pass


# ---------------------------------------------------------------------------
# Materials
# ---------------------------------------------------------------------------
def make_mat(name, color, rough=0.6, sss=0.0):
    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
            if sss > 0.0:
                try:
                    b.inputs["Subsurface Weight"].default_value = float(sss)
                except Exception:
                    pass
    except Exception:
        pass
    return m


MAT = {
    "stalk": make_mat("Stalk", (0.28, 0.45, 0.12), 0.5),
    "sheath": make_mat("Sheath", (0.25, 0.46, 0.13), 0.55),
    "leaf": make_mat("Leaf", (0.16, 0.40, 0.08), 0.5, 0.1),
    "tassel": make_mat("Tassel", (0.62, 0.58, 0.30), 0.7),
    "anther": make_mat("Anther", (0.85, 0.75, 0.25), 0.7),
    "husk": make_mat("Husk", (0.45, 0.62, 0.25), 0.6, 0.1),
    "cob": make_mat("Cob", (0.85, 0.80, 0.70), 0.8),
    "kernel": make_mat("Kernel", (0.95, 0.72, 0.12), 0.35, 0.1),
    "silk": make_mat("Silk", (0.75, 0.55, 0.30), 0.6),
    "brace": make_mat("BraceRoot", (0.40, 0.42, 0.20), 0.6),
    "root": make_mat("Root", (0.82, 0.76, 0.60), 0.8),
}


# ---------------------------------------------------------------------------
# Geometry helpers
# ---------------------------------------------------------------------------
def add_tube(bm, pts, radii, seg=8, cap=True):
    n = len(pts)
    if n < 2:
        return
    tans = []
    for i in range(n):
        if i == 0:
            t = pts[1] - pts[0]
        elif i == n - 1:
            t = pts[-1] - pts[-2]
        else:
            t = pts[i + 1] - pts[i - 1]
        if t.length < 1e-9:
            t = tans[-1].copy() if tans else Vector((0.0, 0.0, 1.0))
        tans.append(t.normalized())
    nrm = tans[0].orthogonal().normalized()
    rings = []
    for i in range(n):
        t = tans[i]
        nrm = nrm - t * nrm.dot(t)
        if nrm.length < 1e-9:
            nrm = t.orthogonal()
        nrm.normalize()
        b = t.cross(nrm)
        ring = []
        r = max(1e-5, radii[i])
        for s in range(seg):
            a = 2.0 * math.pi * s / seg
            ring.append(bm.verts.new(pts[i] + (nrm * math.cos(a) + b * math.sin(a)) * r))
        rings.append(ring)
    for i in range(n - 1):
        for s in range(seg):
            s2 = (s + 1) % seg
            bm.faces.new((rings[i][s], rings[i][s2], rings[i + 1][s2], rings[i + 1][s]))
    if cap and seg >= 3:
        bm.faces.new(list(reversed(rings[0])))
        bm.faces.new(rings[-1])


def add_ellipsoid(bm, c, radial, tang, s_t, s_a, s_r, subdiv=1):
    radial = radial.normalized()
    tang = (tang - radial * tang.dot(radial))
    if tang.length < 1e-9:
        tang = radial.orthogonal()
    tang.normalize()
    axial = radial.cross(tang).normalized()
    nu = 4 + 2 * subdiv
    nv = 2 + subdiv
    top = bm.verts.new(c + axial * s_a)
    bot = bm.verts.new(c - axial * s_a)
    rings = []
    for i in range(1, nv + 1):
        th = math.pi * i / (nv + 1)
        ring = []
        for j in range(nu):
            ph = 2.0 * math.pi * j / nu
            p = c + axial * (s_a * math.cos(th)) + \
                (radial * (math.cos(ph) * s_r) + tang * (math.sin(ph) * s_t)) * math.sin(th)
            ring.append(bm.verts.new(p))
        rings.append(ring)
    for j in range(nu):
        j2 = (j + 1) % nu
        bm.faces.new((top, rings[0][j], rings[0][j2]))
        bm.faces.new((bot, rings[-1][j2], rings[-1][j]))
    for i in range(nv - 1):
        for j in range(nu):
            j2 = (j + 1) % nu
            bm.faces.new((rings[i][j], rings[i + 1][j], rings[i + 1][j2], rings[i][j2]))


def ellipsoid_axis(bm, c, axis, perp, s_along, s_w, s_d, subdiv=1):
    axis = axis.normalized()
    perp = perp - axis * perp.dot(axis)
    if perp.length < 1e-9:
        perp = axis.orthogonal()
    perp.normalize()
    tang = axis.cross(perp)
    add_ellipsoid(bm, c, perp, tang, s_w, s_along, s_d, subdiv)


def make_obj(bm, name, mat):
    try:
        bmesh.ops.recalc_face_normals(bm, faces=bm.faces[:])
    except Exception:
        pass
    me = bpy.data.meshes.new(name)
    bm.to_mesh(me)
    bm.free()
    me.update()
    ob = bpy.data.objects.new(name, me)
    bpy.context.scene.collection.objects.link(ob)
    try:
        me.materials.append(mat)
    except Exception:
        pass
    try:
        for p in me.polygons:
            p.use_smooth = True
    except Exception:
        pass
    return ob


def add_solidify(ob, th):
    try:
        mod = ob.modifiers.new("Solidify", 'SOLIDIFY')
        mod.thickness = th
    except Exception:
        pass


# ---------------------------------------------------------------------------
# Stalk (culm): nodes and internodes
# ---------------------------------------------------------------------------
INTERNODES = [0.035, 0.055, 0.085, 0.115, 0.145, 0.165, 0.175,
              0.18, 0.18, 0.175, 0.165, 0.155, 0.145, 0.135]
NN = len(INTERNODES)
z_nodes = [0.0]
for _L in INTERNODES:
    z_nodes.append(z_nodes[-1] + _L)
Z_TOP = z_nodes[-1]
PED_TOP = Z_TOP + 0.20
TASSEL_LEN = 0.32
STALK_BOTTOM = -0.06
EAR_NODE = 7

leaf_phis = [(i % 2) * math.pi + random.uniform(-0.15, 0.15) for i in range(NN + 1)]


def stalk_radius(z):
    if z < 0.0:
        return 0.0125
    if z <= Z_TOP:
        return 0.0135 - 0.0063 * (z / Z_TOP)
    f = min(1.0, (z - Z_TOP) / (PED_TOP - Z_TOP))
    return 0.0072 - 0.0034 * f


def sheath_r(z):
    return stalk_radius(z) * 1.15 + 0.0008


bm_st = bmesh.new()
pts, radii = [], []
z = STALK_BOTTOM
while z <= PED_TOP + 1e-6:
    r = stalk_radius(z)
    for zn in z_nodes:
        r += 0.0012 * math.exp(-((z - zn) / 0.006) ** 2)
    pts.append(Vector((0.0, 0.0, z)))
    radii.append(r)
    z += 0.005
add_tube(bm_st, pts, radii, seg=16, cap=True)
make_obj(bm_st, "Stalk", MAT["stalk"])


# ---------------------------------------------------------------------------
# Leaves: sheaths wrapping the internodes + arching blades (distichous)
# ---------------------------------------------------------------------------
bm_l = bmesh.new()
bm_sh = bmesh.new()
for i in range(1, NN + 1):
    phi = leaf_phis[i]
    R = Vector((math.cos(phi), math.sin(phi), 0.0))
    S = Vector((-math.sin(phi), math.cos(phi), 0.0))
    zb = z_nodes[i]
    zc = zb + (0.8 * INTERNODES[i] if i < NN else 0.09)

    # sheath
    spts, srad = [], []
    ns = 8
    for k in range(ns + 1):
        zz = zb + (zc - zb) * k / ns
        spts.append(Vector((0.0, 0.0, zz)))
        srad.append(sheath_r(zz) * (1.0 + 0.08 * (k / ns) ** 3))
    add_tube(bm_sh, spts, srad, seg=14, cap=False)

    # blade
    f = i / NN
    L = 0.32 + 0.60 * math.exp(-((f - 0.55) / 0.30) ** 2)
    wmax = 0.035 + 0.06 * math.exp(-((f - 0.5) / 0.35) ** 2)
    th0 = 0.30 + 0.30 * (1.0 - f)
    bend = 1.2 + 0.7 * (1.0 - f) + random.uniform(-0.15, 0.15)
    drift = random.uniform(-0.12, 0.12)
    nu, nv = 40, 8
    p = Vector((0.0, 0.0, zc)) + R * sheath_r(zc)
    rows = []
    for k in range(nu + 1):
        s = k / nu
        th = th0 + bend * s ** 1.6
        d = (Z * math.cos(th) + R * math.sin(th) + S * (drift * s)).normalized()
        n_ = d.cross(S).normalized()
        sd = n_.cross(d).normalized()
        if k > 0:
            p = p + d * (L / nu)
        if p.z < 0.01:
            p = Vector((p.x, p.y, 0.01))
        w = wmax * min(1.0, 0.4 + 2.2 * s) * (1.0 - s) ** 0.75
        row = []
        for j in range(nv + 1):
            v = 2.0 * j / nv - 1.0
            wave = 0.0035 * math.sin(s * L * 60.0 + i) * abs(v) ** 2 * (w / wmax)
            q = p + sd * (v * w * 0.5) + n_ * (abs(v) * w * 0.12 + wave)
            row.append(bm_l.verts.new(q))
        rows.append(row)
    for k in range(nu):
        for j in range(nv):
            bm_l.faces.new((rows[k][j], rows[k][j + 1], rows[k + 1][j + 1], rows[k + 1][j]))

make_obj(bm_sh, "Leaf_Sheaths", MAT["sheath"])
ob_leaves = make_obj(bm_l, "Leaf_Blades", MAT["leaf"])
add_solidify(ob_leaves, 0.0006)


# ---------------------------------------------------------------------------
# Tassel (male inflorescence): central spike + lateral branches + spikelets
# ---------------------------------------------------------------------------
bm_t = bmesh.new()
bm_an = bmesh.new()


def add_spikelets(pts, j0, per, size):
    for j in range(j0, len(pts) - 1):
        T = pts[j + 1] - pts[j]
        if T.length < 1e-9:
            continue
        T.normalize()
        base_perp = T.orthogonal().normalized()
        for q in range(per):
            ang = 2.0 * math.pi * q / per + j * 0.9
            perp = (Matrix.Rotation(ang, 3, T) @ base_perp).normalized()
            dirn = (T * 0.85 + perp * 0.5).normalized()
            c = pts[j] + perp * (size * 0.9) + dirn * (size * 0.6)
            ellipsoid_axis(bm_t, c, dirn, perp, size * 1.6, size * 0.55, size * 0.45, 1)
            if random.random() < 0.35:
                a = random.uniform(0.0, 2.0 * math.pi)
                ap = Vector((math.cos(a), math.sin(a), 0.0))
                ac = c + perp * (size * 0.5) - Z * (size * 1.5)
                ellipsoid_axis(bm_an, ac, Z, ap, size * 0.9, size * 0.25, size * 0.25, 1)


cpts = []
nC = 24
for j in range(nC + 1):
    s = j / nC
    cpts.append(Vector((0.01 * s * s, 0.0, PED_TOP + TASSEL_LEN * s)))
crad = [0.0036 * (1.0 - 0.65 * j / nC) for j in range(nC + 1)]
add_tube(bm_t, cpts, crad, seg=6)
add_spikelets(cpts, 4, 4, 0.0025)

nb = 12
for b in range(nb):
    hb = PED_TOP + 0.10 * (b / nb)
    az = b * 2.39996
    Rb = Vector((math.cos(az), math.sin(az), 0.0))
    L = 0.22 - 0.08 * (b / nb) + random.uniform(-0.02, 0.02)
    n = max(8, int(L / 0.01))
    p = Vector((0.01 * ((hb - PED_TOP) / TASSEL_LEN) ** 2, 0.0, hb)) + Rb * 0.002
    bpts = [p.copy()]
    for j in range(n):
        s = (j + 1) / n
        th = 0.55 + 1.1 * s ** 1.5
        d = Z * math.cos(th) + Rb * math.sin(th)
        p = p + d * (L / n)
        bpts.append(p.copy())
    brad = [0.0018 * (1.0 - 0.6 * j / n) for j in range(n + 1)]
    add_tube(bm_t, bpts, brad, seg=5)
    add_spikelets(bpts, max(1, n // 6), 2, 0.0025)

make_obj(bm_t, "Tassel", MAT["tassel"])
make_obj(bm_an, "Tassel_Anthers", MAT["anther"])


# ---------------------------------------------------------------------------
# Ears (female inflorescences): shank, husks, cob, kernels, silks
# ---------------------------------------------------------------------------
def build_ear(node, phi, scale, tag):
    zn = z_nodes[node]
    R = Vector((math.cos(phi), math.sin(phi), 0.0))
    base = Vector((0.0, 0.0, zn)) + R * (stalk_radius(zn) * 0.9)
    tilt = math.radians(31.0 if scale > 0.8 else 36.0)
    axis = (Z * math.cos(tilt) + R * math.sin(tilt)).normalized()
    xw = (R * math.cos(tilt) - Z * math.sin(tilt)).normalized()
    yw = axis.cross(xw).normalized()
    M3 = Matrix((xw, yw, axis)).transposed()

    def W(v):
        return base + M3 @ v

    Ls = 0.07 * scale
    Le = 0.19 * scale
    tipL = 0.04 * scale

    def ear_r(t):
        t = max(0.0, min(1.0, t))
        return scale * 0.022 * (1.0 - 0.45 * t ** 1.6) * min(1.0, 0.75 + 2.0 * t)

    def kdep(t):
        return ear_r(t) * 0.42

    def husk_r(zl):
        t = (zl - Ls) / Le
        if t < 0.0:
            return ear_r(0.0) * 0.9
        if t <= 1.0:
            return ear_r(t)
        f = min(1.0, (zl - Ls - Le) / tipL)
        return ear_r(1.0) * (1.0 - f) + 0.0025 * scale * f

    # shank
    bm_sk = bmesh.new()
    spts = [W(Vector((0.0, 0.0, -0.005 + (Ls + 0.01) * j / 8))) for j in range(9)]
    add_tube(bm_sk, spts, [0.006 * scale] * 9, seg=10)
    make_obj(bm_sk, "Ear_%s_Shank" % tag, MAT["stalk"])

    # husks
    bm_h = bmesh.new()
    nh = 7
    nu, nv = 28, 8
    z0 = Ls - 0.005 * scale
    z1 = Ls + Le + tipL
    for h in range(nh):
        az0 = 2.0 * math.pi * h / nh + random.uniform(-0.2, 0.2)
        arc = random.uniform(2.2, 2.8)
        off = (0.0015 + 0.0011 * h) * scale
        grid = []
        for i in range(nu + 1):
            u = i / nu
            zl = z0 + (z1 - z0) * u
            r = husk_r(zl) + off * (1.0 - 0.7 * u)
            a_w = arc * (1.0 - 0.55 * u ** 2)
            row = []
            for j in range(nv + 1):
                a = az0 + a_w * (j / nv - 0.5)
                row.append(bm_h.verts.new(W(Vector((r * math.cos(a), r * math.sin(a), zl)))))
            grid.append(row)
        for i in range(nu):
            for j in range(nv):
                bm_h.faces.new((grid[i][j], grid[i][j + 1], grid[i + 1][j + 1], grid[i + 1][j]))
    ob_h = make_obj(bm_h, "Ear_%s_Husks" % tag, MAT["husk"])
    add_solidify(ob_h, 0.0005)

    # cob
    bm_c = bmesh.new()
    n = 24
    cpts_ = [W(Vector((0.0, 0.0, Ls + Le * j / n))) for j in range(n + 1)]
    crad_ = [max(0.002, ear_r(j / n) - kdep(j / n)) for j in range(n + 1)]
    add_tube(bm_c, cpts_, crad_, seg=12)
    make_obj(bm_c, "Ear_%s_Cob" % tag, MAT["cob"])

    # kernels in paired longitudinal rows
    bm_k = bmesh.new()
    nrows = 16 if scale > 0.8 else 12
    nk = max(4, int(Le * 0.93 / 0.0056))
    for k in range(nk):
        t = 0.03 + 0.92 * (k + 0.5) / nk
        er = ear_r(t)
        kd = kdep(t)
        rc = er - kd
        zl = Ls + Le * t
        for r in range(nrows):
            az = 2.0 * math.pi * r / nrows + 0.03 * math.sin(k * 0.4)
            rho = Vector((math.cos(az), math.sin(az), 0.0))
            tau = Vector((-math.sin(az), math.cos(az), 0.0))
            c = W(rho * (rc + kd * 0.5) + Vector((0.0, 0.0, zl)))
            sx = math.pi * er / nrows * 0.92
            sy = Le * 0.93 / nk * 0.5 * 0.95
            sz = kd * 0.5
            add_ellipsoid(bm_k, c, M3 @ rho, M3 @ tau, sx, sy, sz, subdiv=2)
    make_obj(bm_k, "Ear_%s_Kernels" % tag, MAT["kernel"])

    # silks emerging from husk tip
    bm_s = bmesh.new()
    zs = Ls + Le + tipL * 0.8
    for qd in range(70):
        a = random.uniform(0.0, 2.0 * math.pi)
        rr = random.uniform(0.0, 0.0025 * scale)
        st = W(Vector((rr * math.cos(a), rr * math.sin(a), zs)))
        spread = random.uniform(0.2, 0.6)
        dv = (M3 @ Vector((math.cos(a) * spread, math.sin(a) * spread, 1.0))).normalized()
        Ls_ = random.uniform(0.06, 0.14) * scale
        n = 9
        pts_ = [st.copy()]
        p = st.copy()
        for k in range(n):
            dv = (dv + Vector((0.0, 0.0, -0.22)) +
                  Vector((random.uniform(-0.08, 0.08),
                          random.uniform(-0.08, 0.08),
                          random.uniform(-0.04, 0.04)))).normalized()
            p = p + dv * (Ls_ / n)
            pts_.append(p.copy())
        radii_ = [0.00035 * (1.0 - 0.5 * k / n) for k in range(n + 1)]
        add_tube(bm_s, pts_, radii_, seg=4, cap=False)
    make_obj(bm_s, "Ear_%s_Silks" % tag, MAT["silk"])


try:
    build_ear(EAR_NODE, leaf_phis[EAR_NODE] + 0.45, 1.0, "Primary")
except Exception as e:
    print("Primary ear failed:", e)

try:
    build_ear(EAR_NODE - 1, leaf_phis[EAR_NODE - 1] + 0.45, 0.6, "Secondary")
except Exception as e:
    print("Secondary ear failed:", e)


# ---------------------------------------------------------------------------
# Brace (prop) roots from lower above-ground nodes
# ---------------------------------------------------------------------------
bm_b = bmesh.new()
for idx in (1, 2):
    zn = z_nodes[idx]
    nroot = 8 if idx == 1 else 10
    for k in range(nroot):
        az = 2.0 * math.pi * k / nroot + random.uniform(-0.2, 0.2) + idx * 0.3
        R = Vector((math.cos(az), math.sin(az), 0.0))
        r0 = stalk_radius(zn) * 0.95
        D = 0.07 + 0.05 * (idx - 1) + random.uniform(0.0, 0.03)
        P0 = Vector((0.0, 0.0, zn + random.uniform(-0.003, 0.003))) + R * r0
        P2 = R * D + Vector((0.0, 0.0, -0.035))
        P1 = R * (r0 + 0.6 * (D - r0)) + Vector((0.0, 0.0, zn * 0.85))
        n = 16
        pts, radii = [], []
        for j in range(n + 1):
            t = j / n
            p = P0 * ((1 - t) ** 2) + P1 * (2 * (1 - t) * t) + P2 * (t * t)
            pts.append(p)
            radii.append(0.0032 * (1.0 - 0.3 * t))
        add_tube(bm_b, pts, radii, seg=8)
make_obj(bm_b, "Brace_Roots", MAT["brace"])


# ---------------------------------------------------------------------------
# Underground fibrous root system: crown roots, seminal root, laterals
# ---------------------------------------------------------------------------
def root_walk(start, d, L, step, grav, jit=0.25):
    pts = [start.copy()]
    d = d.normalized()
    p = start.copy()
    for _i in range(max(2, int(L / step))):
        d = d + Vector((random.uniform(-jit, jit),
                        random.uniform(-jit, jit),
                        random.uniform(-jit, jit) * 0.5)) - Z * grav
        if p.z > -0.01 and d.z > 0.0:
            d.z = -abs(d.z)
        if d.length < 1e-6:
            d = Vector((0.0, 0.0, -1.0))
        d.normalize()
        p = p + d * step
        if p.z > -0.005:
            p.z = -0.005
        pts.append(p.copy())
    return pts


def add_laterals(bm, pts, count, r_lat):
    n = len(pts)
    if n < 4:
        return
    for _l in range(count):
        idx = random.randint(1, n - 3)
        tang = pts[idx + 1] - pts[idx]
        if tang.length < 1e-9:
            continue
        tang.normalize()
        rv = Vector((random.uniform(-1, 1), random.uniform(-1, 1), random.uniform(-1, 1)))
        dl = rv - tang * rv.dot(tang)
        if dl.length < 1e-6:
            dl = tang.orthogonal()
        dl = (dl.normalized() - Z * 0.2).normalized()
        lp = root_walk(pts[idx], dl, random.uniform(0.03, 0.1), 0.01, 0.03, 0.3)
        lr = [r_lat * (1.0 - 0.6 * j / (len(lp) - 1)) for j in range(len(lp))]
        add_tube(bm, lp, lr, seg=4)


bm_r = bmesh.new()

# crown (nodal) roots from underground nodes
for k in range(24):
    az = 2.0 * math.pi * k / 24 + random.uniform(-0.12, 0.12)
    zst = random.uniform(-0.055, -0.01)
    R = Vector((math.cos(az), math.sin(az), 0.0))
    start = Vector((0.0, 0.0, zst)) + R * stalk_radius(zst) * 0.9
    dip = math.radians(random.uniform(15.0, 55.0))
    d = R * math.cos(dip) - Z * math.sin(dip)
    L = random.uniform(0.25, 0.55)
    pts = root_walk(start, d, L, 0.02, 0.06)
    radii = [0.0022 * (1.0 - 0.6 * j / (len(pts) - 1)) for j in range(len(pts))]
    add_tube(bm_r, pts, radii, seg=6)
    add_laterals(bm_r, pts, random.randint(4, 7), 0.0006)

# primary (seminal) root
start = Vector((0.0, 0.0, STALK_BOTTOM + 0.002))
pts = root_walk(start, Vector((0.05, 0.0, -1.0)), 0.5, 0.02, 0.08, 0.12)
radii = [0.0018 * (1.0 - 0.6 * j / (len(pts) - 1)) for j in range(len(pts))]
add_tube(bm_r, pts, radii, seg=6)
add_laterals(bm_r, pts, 8, 0.0006)

# a few seminal roots from the scutellar node
for k in range(4):
    az = 2.0 * math.pi * k / 4 + 0.4
    R = Vector((math.cos(az), math.sin(az), 0.0))
    d = R * 0.6 - Z * 0.8
    pts = root_walk(start + R * 0.004, d, random.uniform(0.25, 0.4), 0.02, 0.07)
    radii = [0.0014 * (1.0 - 0.6 * j / (len(pts) - 1)) for j in range(len(pts))]
    add_tube(bm_r, pts, radii, seg=5)
    add_laterals(bm_r, pts, 4, 0.0005)

make_obj(bm_r, "Roots", MAT["root"])

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